diff --git a/metadata-ingestion/scripts/modeldocgen.py b/metadata-ingestion/scripts/modeldocgen.py index ee5f06cb801ba..998947e5caa95 100644 --- a/metadata-ingestion/scripts/modeldocgen.py +++ b/metadata-ingestion/scripts/modeldocgen.py @@ -14,7 +14,11 @@ import click from datahub.configuration.common import ConfigEnum, PermissiveConfigModel -from datahub.emitter.mce_builder import make_data_platform_urn, make_dataset_urn +from datahub.emitter.mce_builder import ( + make_data_platform_urn, + make_dataset_urn, + make_schema_field_urn, +) from datahub.emitter.mcp import MetadataChangeProposalWrapper from datahub.emitter.rest_emitter import DatahubRestEmitter from datahub.ingestion.api.common import PipelineContext, RecordEnvelope @@ -442,10 +446,10 @@ def strip_types(field_path: str) -> str: name=relnship_name, foreignDataset=foreign_dataset_urn, foreignFields=[ - f"urn:li:schemaField:({foreign_dataset_urn}, urn)" + make_schema_field_urn(foreign_dataset_urn, "urn") ], sourceFields=[ - f"urn:li:schemaField:({source_dataset_urn},{f_field.fieldPath})" + make_schema_field_urn(source_dataset_urn, f_field.fieldPath) ], ) foreign_keys.append(fkey) diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py index 238fd88f1c950..e5779791ed412 100644 --- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py +++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py @@ -32,6 +32,7 @@ make_data_platform_urn, make_dataplatform_instance_urn, make_dataset_urn_with_platform_instance, + make_schema_field_urn, make_tag_urn, ) from datahub.emitter.mcp import MetadataChangeProposalWrapper @@ -669,7 +670,7 @@ def get_foreign_key_metadata( ) source_fields = [ - f"urn:li:schemaField:({dataset_urn},{f})" + make_schema_field_urn(dataset_urn, f) for f in fk_dict["constrained_columns"] ] foreign_dataset = make_dataset_urn_with_platform_instance( @@ -679,7 +680,7 @@ def get_foreign_key_metadata( env=self.config.env, ) foreign_fields = [ - f"urn:li:schemaField:({foreign_dataset},{f})" + make_schema_field_urn(foreign_dataset, f) for f in fk_dict["referred_columns"] ] diff --git a/metadata-ingestion/tests/unit/test_sql_common.py b/metadata-ingestion/tests/unit/test_sql_common.py index a98bf64171122..cfb8f55bd977f 100644 --- a/metadata-ingestion/tests/unit/test_sql_common.py +++ b/metadata-ingestion/tests/unit/test_sql_common.py @@ -38,7 +38,7 @@ def test_generate_foreign_key(): "referred_columns": ["test_referred_column"], # type: ignore } foreign_key = source.get_foreign_key_metadata( - dataset_urn="test_urn", + dataset_urn="urn:li:dataset:(urn:li:dataPlatform:TEST,test_schema.base_urn,PROD)", schema="test_schema", fk_dict=fk_dict, inspector=mock.Mock(), @@ -48,7 +48,9 @@ def test_generate_foreign_key(): assert [ "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:TEST,test_referred_schema.test_table,PROD),test_referred_column)" ] == foreign_key.foreignFields - assert ["urn:li:schemaField:(test_urn,test_column)"] == foreign_key.sourceFields + assert [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:TEST,test_schema.base_urn,PROD),test_column)" + ] == foreign_key.sourceFields def test_use_source_schema_for_foreign_key_if_not_specified(): @@ -60,7 +62,7 @@ def test_use_source_schema_for_foreign_key_if_not_specified(): "referred_columns": ["test_referred_column"], # type: ignore } foreign_key = source.get_foreign_key_metadata( - dataset_urn="test_urn", + dataset_urn="urn:li:dataset:(urn:li:dataPlatform:TEST,test_schema.base_urn,PROD)", schema="test_schema", fk_dict=fk_dict, inspector=mock.Mock(), @@ -70,7 +72,9 @@ def test_use_source_schema_for_foreign_key_if_not_specified(): assert [ "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:TEST,test_schema.test_table,PROD),test_referred_column)" ] == foreign_key.foreignFields - assert ["urn:li:schemaField:(test_urn,test_column)"] == foreign_key.sourceFields + assert [ + "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:TEST,test_schema.base_urn,PROD),test_column)" + ] == foreign_key.sourceFields PLATFORM_FROM_SQLALCHEMY_URI_TEST_CASES: Dict[str, str] = {