Merge branch 'master' into datahub-source-postgres

datahub-project · Aug 30, 2023 · e77e84e · e77e84e
2 parents ec6cbb4 + 4c69f9a
commit e77e84e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py b/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py
@@ -9,7 +9,7 @@
 from sqlalchemy.engine.reflection import Inspector
 
 from datahub.configuration.validate_field_rename import pydantic_renamed_field
-from datahub.emitter.mcp_builder import ContainerKey
+from datahub.emitter.mcp_builder import ContainerKey, DatabaseKey
 from datahub.ingestion.api.decorators import (
     SourceCapability,
     SupportStatus,
@@ -205,7 +205,7 @@ def gen_schema_containers(
             extra_properties=extra_properties,
         )
 
-    def get_database_container_key(self, db_name: str, schema: str) -> ContainerKey:
+    def get_database_container_key(self, db_name: str, schema: str) -> DatabaseKey:
         # Because our overridden get_allowed_schemas method returns db_name as the schema name,
         # the db_name and schema here will be the same. Hence, we just ignore the schema parameter.
         # Based on community feedback, db_name only available if it is explicitly specified in the connection string.

diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_utils.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_utils.py
@@ -35,7 +35,7 @@ def gen_schema_key(
     platform: str,
     platform_instance: Optional[str],
     env: Optional[str],
-) -> ContainerKey:
+) -> SchemaKey:
     return SchemaKey(
         database=db_name,
         schema=schema,
@@ -48,7 +48,7 @@ def gen_schema_key(
 
 def gen_database_key(
     database: str, platform: str, platform_instance: Optional[str], env: Optional[str]
-) -> ContainerKey:
+) -> DatabaseKey:
     return DatabaseKey(
         database=database,
         platform=platform,