""" OpenMetadata high-level API Database test """ import uuid from unittest import TestCase from metadata.generated.schema.api.data.createDatabase import ( CreateDatabaseEntityRequest, ) from metadata.generated.schema.api.services.createDatabaseService import ( CreateDatabaseServiceEntityRequest, ) from metadata.generated.schema.api.teams.createUser import CreateUserEntityRequest from metadata.generated.schema.entity.data.database import Database from metadata.generated.schema.entity.services.databaseService import ( DatabaseService, DatabaseServiceType, ) from metadata.generated.schema.type.entityReference import EntityReference from metadata.generated.schema.type.jdbcConnection import JdbcInfo from metadata.ingestion.ometa.ometa_api import OpenMetadata from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig class OMetaDatabaseTest(TestCase): """ Run this integration test with the local API available Install the ingestion package before running the tests """ service_entity_id = None server_config = MetadataServerConfig(api_endpoint="http://localhost:8585/api") metadata = OpenMetadata(server_config) user = metadata.create_or_update( data=CreateUserEntityRequest(name="random-user", email="random@user.com"), ) owner = EntityReference(id=user.id, type="user") service = CreateDatabaseServiceEntityRequest( name="test-service-db", serviceType=DatabaseServiceType.MySQL, jdbc=JdbcInfo(driverClass="jdbc", connectionUrl="jdbc://localhost"), ) @classmethod def setUpClass(cls) -> None: """ Prepare ingredients """ cls.service_entity = cls.metadata.create_or_update(data=cls.service) cls.entity = Database( id=uuid.uuid4(), name="test-db", service=EntityReference(id=cls.service_entity.id, type="databaseService"), fullyQualifiedName="test-service-db.test-db", ) cls.create = CreateDatabaseEntityRequest( name="test-db", service=EntityReference(id=cls.service_entity.id, type="databaseService"), ) @classmethod def tearDownClass(cls) -> None: """ Clean up """ db_id = str( cls.metadata.get_by_name( entity=Database, fqdn="test-service-db.test-db" ).id.__root__ ) service_id = str( cls.metadata.get_by_name( entity=DatabaseService, fqdn="test-service-db" ).id.__root__ ) cls.metadata.delete(entity=Database, entity_id=db_id) cls.metadata.delete(entity=DatabaseService, entity_id=service_id) def test_create(self): """ We can create a Database 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 = CreateDatabaseEntityRequest(**updated) res = self.metadata.create_or_update(data=updated_entity) # Same ID, updated algorithm self.assertEqual(res.service.id, updated_entity.service.id) self.assertEqual(res_create.id, res.id) self.assertEqual(res.owner.id, self.user.id) def test_get_name(self): """ We can fetch a Database by name and get it back as Entity """ self.metadata.create_or_update(data=self.create) res = self.metadata.get_by_name( entity=Database, fqdn=self.entity.fullyQualifiedName ) self.assertEqual(res.name, self.entity.name) def test_get_id(self): """ We can fetch a Database 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=Database, fqdn=self.entity.fullyQualifiedName ) # Then fetch by ID res = self.metadata.get_by_id( entity=Database, entity_id=str(res_name.id.__root__) ) self.assertEqual(res_name.id, res.id) def test_list(self): """ We can list all our Database """ self.metadata.create_or_update(data=self.create) res = self.metadata.list_entities(entity=Database) # 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 Database by ID """ self.metadata.create_or_update(data=self.create) # Find by name res_name = self.metadata.get_by_name( entity=Database, fqdn=self.entity.fullyQualifiedName ) # Then fetch by ID res_id = self.metadata.get_by_id( entity=Database, entity_id=str(res_name.id.__root__) ) # Delete self.metadata.delete(entity=Database, entity_id=str(res_id.id.__root__)) # Then we should not find it res = self.metadata.list_entities(entity=Database) assert not next( iter(ent for ent in res.entities if ent.name == self.entity.name), None )