diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py b/metadata-ingestion/src/datahub/ingestion/source/sql/athena.py index 53501b9a536ee..9cb613bde1e9f 100644 --- 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 index c5baf148b0e5e..723a8c5fd8669 100644 --- 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,