Add service type to container (#10441)

ingestion/tests/integration/ometa/test_ometa_objectstore_api.py

@@ -0,0 +1,263 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+OpenMetadata high-level API Container test
+"""
+import uuid
+from unittest import TestCase
+
+from metadata.generated.schema.api.data.createContainer import CreateContainerRequest
+from metadata.generated.schema.api.services.createObjectStoreService import (
+    CreateObjectStoreServiceRequest,
+)
+from metadata.generated.schema.api.teams.createUser import CreateUserRequest
+from metadata.generated.schema.entity.data.container import Container
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.entity.services.connections.objectstore.s3ObjectStoreConnection import (
+    S3StoreConnection,
+)
+from metadata.generated.schema.entity.services.objectstoreService import (
+    ObjectStoreConnection,
+    ObjectStoreService,
+    ObjectStoreServiceType,
+)
+from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
+    OpenMetadataJWTClientConfig,
+)
+from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
+from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+
+
+class OMetaObjectStoreTest(TestCase):
+    """
+    Run this integration test with the local API available
+    Install the ingestion package before running the tests
+    """
+
+    service_entity_id = None
+
+    server_config = OpenMetadataConnection(
+        hostPort="http://localhost:8585/api",
+        authProvider="openmetadata",
+        securityConfig=OpenMetadataJWTClientConfig(
+            jwtToken="eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"
+        ),
+    )
+    metadata = OpenMetadata(server_config)
+
+    assert metadata.health_check()
+
+    user = metadata.create_or_update(
+        data=CreateUserRequest(name="random-user", email="random@user.com"),
+    )
+    owner = EntityReference(id=user.id, type="user")
+
+    service = CreateObjectStoreServiceRequest(
+        name="test-service-object",
+        serviceType=ObjectStoreServiceType.S3,
+        connection=ObjectStoreConnection(
+            config=S3StoreConnection(awsConfig=AWSCredentials(awsRegion="us-east-2"))
+        ),
+    )
+    service_type = "objectstoreService"
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        """
+        Prepare ingredients
+        """
+        cls.service_entity = cls.metadata.create_or_update(data=cls.service)
+
+        cls.entity = Container(
+            id=uuid.uuid4(),
+            name="test",
+            service=EntityReference(id=cls.service_entity.id, type=cls.service_type),
+            fullyQualifiedName="test-service-object.test",
+        )
+
+        cls.create = CreateContainerRequest(
+            name="test",
+            service=cls.service_entity.fullyQualifiedName,
+        )
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        """
+        Clean up
+        """
+
+        service_id = str(
+            cls.metadata.get_by_name(
+                entity=ObjectStoreService, fqn="test-service-object"
+            ).id.__root__
+        )
+
+        cls.metadata.delete(
+            entity=ObjectStoreService,
+            entity_id=service_id,
+            recursive=True,
+            hard_delete=True,
+        )
+
+    def test_create(self):
+        """
+        We can create a Container and we receive it back as Entity
+        """
+
+        res = self.metadata.create_or_update(data=self.create)
+
+        self.assertEqual(res.name, self.entity.name)
+        self.assertEqual(res.service.id, self.entity.service.id)
+        self.assertEqual(res.owner, None)
+
+    def test_update(self):
+        """
+        Updating it properly changes its properties
+        """
+
+        res_create = self.metadata.create_or_update(data=self.create)
+
+        updated = self.create.dict(exclude_unset=True)
+        updated["owner"] = self.owner
+        updated_entity = CreateContainerRequest(**updated)
+
+        res = self.metadata.create_or_update(data=updated_entity)
+
+        # Same ID, updated algorithm
+        self.assertEqual(
+            res.service.fullyQualifiedName, updated_entity.service.__root__
+        )
+        self.assertEqual(res_create.id, res.id)
+        self.assertEqual(res.owner.id, self.user.id)
+
+    def test_get_name(self):
+        """
+        We can fetch a Container by name and get it back as Entity
+        """
+
+        self.metadata.create_or_update(data=self.create)
+
+        res = self.metadata.get_by_name(
+            entity=Container, fqn=self.entity.fullyQualifiedName
+        )
+        self.assertEqual(res.name, self.entity.name)
+
+    def test_get_id(self):
+        """
+        We can fetch a Container by ID and get it back as Entity
+        """
+
+        self.metadata.create_or_update(data=self.create)
+
+        # First pick up by name
+        res_name = self.metadata.get_by_name(
+            entity=Container, fqn=self.entity.fullyQualifiedName
+        )
+        # Then fetch by ID
+        res = self.metadata.get_by_id(entity=Container, entity_id=res_name.id)
+
+        self.assertEqual(res_name.id, res.id)
+
+    def test_list(self):
+        """
+        We can list all our Containers
+        """
+
+        self.metadata.create_or_update(data=self.create)
+
+        res = self.metadata.list_entities(entity=Container, limit=100)
+
+        # Fetch our test Database. We have already inserted it, so we should find it
+        data = next(
+            iter(ent for ent in res.entities if ent.name == self.entity.name), None
+        )
+        assert data
+
+    def test_delete(self):
+        """
+        We can delete a Container by ID
+        """
+
+        self.metadata.create_or_update(data=self.create)
+
+        # Find by name
+        res_name = self.metadata.get_by_name(
+            entity=Container, fqn=self.entity.fullyQualifiedName
+        )
+        # Then fetch by ID
+        res_id = self.metadata.get_by_id(
+            entity=Container, entity_id=str(res_name.id.__root__)
+        )
+
+        # Delete
+        self.metadata.delete(
+            entity=Container, entity_id=str(res_id.id.__root__), recursive=True
+        )
+
+        # Then we should not find it
+        res = self.metadata.list_entities(entity=Container)
+        assert not next(
+            iter(
+                ent
+                for ent in res.entities
+                if ent.fullyQualifiedName == self.entity.fullyQualifiedName
+            ),
+            None,
+        )
+
+    def test_list_versions(self):
+        """
+        test list Container entity versions
+        """
+        self.metadata.create_or_update(data=self.create)
+
+        # Find by name
+        res_name = self.metadata.get_by_name(
+            entity=Container, fqn=self.entity.fullyQualifiedName
+        )
+
+        res = self.metadata.get_list_entity_versions(
+            entity=Container, entity_id=res_name.id.__root__
+        )
+        assert res
+
+    def test_get_entity_version(self):
+        """
+        test get Container entity version
+        """
+        self.metadata.create_or_update(data=self.create)
+
+        # Find by name
+        res_name = self.metadata.get_by_name(
+            entity=Container, fqn=self.entity.fullyQualifiedName
+        )
+        res = self.metadata.get_entity_version(
+            entity=Container, entity_id=res_name.id.__root__, version=0.1
+        )
+
+        # check we get the correct version requested and the correct entity ID
+        assert res.version.__root__ == 0.1
+        assert res.id == res_name.id
+
+    def test_get_entity_ref(self):
+        """
+        test get EntityReference
+        """
+        res = self.metadata.create_or_update(data=self.create)
+        entity_ref = self.metadata.get_entity_reference(
+            entity=Container, fqn=res.fullyQualifiedName
+        )
+
+        assert res.id == entity_ref.id

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/ContainerRepository.java

@@ -109,7 +109,8 @@ public class ContainerRepository extends EntityRepository<Container> {
   public void prepare(Container container) throws IOException {
     // the objectStoreService is not fully filled in terms of props - go to the db and get it in full and re-set it
     ObjectStoreService objectStoreService = Entity.getEntity(container.getService(), "", Include.NON_DELETED);
-    container.withService(objectStoreService.getEntityReference());
+    container.setService(objectStoreService.getEntityReference());
+    container.setServiceType(objectStoreService.getServiceType());
 
     if (container.getParent() != null) {
       Container parent = Entity.getEntity(container.getParent(), "owner", ALL);

openmetadata-spec/src/main/resources/json/schema/entity/data/container.json

@@ -132,6 +132,10 @@
       },
       "default": null
     },
+    "serviceType": {
+      "description": "Service type this table is hosted in.",
+      "$ref": "../services/objectstoreService.json#/definitions/objectStoreServiceType"
+    },
     "followers": {
       "description": "Followers of this container.",
       "$ref": "../../type/entityReference.json#/definitions/entityReferenceList"