acryl-datahub 1.0.0rc13__py3-none-any.whl → 1.0.0rc14__py3-none-any.whl

datahub/ingestion/sink/datahub_rest.py

@@ -20,7 +20,9 @@ from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
     BATCH_INGEST_MAX_PAYLOAD_LENGTH,
+    DEFAULT_REST_SINK_ENDPOINT,
     DataHubRestEmitter,
+    RestSinkEndpoint,
 )
 from datahub.ingestion.api.common import RecordEnvelope, WorkUnit
 from datahub.ingestion.api.sink import (
@@ -66,6 +68,7 @@ _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
 
 class DatahubRestSinkConfig(DatahubClientConfig):
     mode: RestSinkMode = _DEFAULT_REST_SINK_MODE
+    endpoint: RestSinkEndpoint = DEFAULT_REST_SINK_ENDPOINT
 
     # These only apply in async modes.
     max_threads: pydantic.PositiveInt = _DEFAULT_REST_SINK_MAX_THREADS
@@ -172,6 +175,7 @@ class DatahubRestSink(Sink[DatahubRestSinkConfig, DataHubRestSinkReport]):
             ca_certificate_path=config.ca_certificate_path,
             client_certificate_path=config.client_certificate_path,
             disable_ssl_verification=config.disable_ssl_verification,
+            openapi_ingestion=config.endpoint == RestSinkEndpoint.OPENAPI,
         )
 
     @property

datahub/ingestion/source/common/subtypes.py

@@ -92,3 +92,8 @@ class BIAssetSubTypes(StrEnum):
     # SAP Analytics Cloud
     SAC_STORY = "Story"
     SAC_APPLICATION = "Application"
+
+
+class MLAssetSubTypes(StrEnum):
+    MLFLOW_TRAINING_RUN = "ML Training Run"
+    MLFLOW_EXPERIMENT = "ML Experiment"

datahub/ingestion/source/data_lake_common/path_spec.py

@@ -454,10 +454,8 @@ class PathSpec(ConfigModel):
                     return None
                 partition = partition_split[0]
                 # If partition is in the form of /value1/value2/value3 we infer it from the path and assign partition_0, partition_1, partition_2 etc
-                num = 0
-                for partition_value in partition.split("/"):
+                for num, partition_value in enumerate(partition.split("/")):
                     partition_keys.append((f"partition_{num}", partition_value))
-                    num += 1
             return partition_keys
 
         return None

datahub/ingestion/source/dbt/dbt_common.py

@@ -1774,10 +1774,8 @@ class DBTSourceBase(StatefulIngestionSourceBase):
                     logger.debug(
                         f"Owner after applying owner extraction pattern:'{self.config.owner_extraction_pattern}' is '{owner}'."
                     )
-            if isinstance(owner, list):
-                owners = owner
-            else:
-                owners = [owner]
+            owners = owner if isinstance(owner, list) else [owner]
+
             for owner in owners:
                 if self.config.strip_user_ids_from_email:
                     owner = owner.split("@")[0]

datahub/ingestion/source/dbt/dbt_tests.py

@@ -57,15 +57,11 @@ def _get_name_for_relationship_test(kw_args: Dict[str, str]) -> Optional[str]:
         # base assertions are violated, bail early
         return None
     m = re.match(r"^ref\(\'(.*)\'\)$", destination_ref)
-    if m:
-        destination_table = m.group(1)
-    else:
-        destination_table = destination_ref
+    destination_table = m.group(1) if m else destination_ref
+
     m = re.search(r"ref\(\'(.*)\'\)", source_ref)
-    if m:
-        source_table = m.group(1)
-    else:
-        source_table = source_ref
+    source_table = m.group(1) if m else source_ref
+
     return f"{source_table}.{column_name} referential integrity to {destination_table}.{dest_field_name}"
 
 

datahub/ingestion/source/dremio/dremio_api.py

@@ -683,11 +683,7 @@ class DremioAPIOperations:
                 # Add end anchor for exact matching
                 regex_pattern = regex_pattern + "$"
 
-        for path in paths:
-            if re.match(regex_pattern, path, re.IGNORECASE):
-                return True
-
-        return False
+        return any(re.match(regex_pattern, path, re.IGNORECASE) for path in paths)
 
     def should_include_container(self, path: List[str], name: str) -> bool:
         """

datahub/ingestion/source/dremio/dremio_aspects.py

@@ -116,10 +116,7 @@ class SchemaFieldTypeMapper:
             data_type = data_type.lower()
             type_class = cls.FIELD_TYPE_MAPPING.get(data_type, NullTypeClass)
 
-            if data_size:
-                native_data_type = f"{data_type}({data_size})"
-            else:
-                native_data_type = data_type
+            native_data_type = f"{data_type}({data_size})" if data_size else data_type
 
         try:
             schema_field_type = SchemaFieldDataTypeClass(type=type_class())

datahub/ingestion/source/dynamodb/dynamodb.py

@@ -246,6 +246,7 @@ class DynamoDBSource(StatefulIngestionSourceBase):
             platform=self.platform,
             platform_instance=platform_instance,
             name=dataset_name,
+            env=self.config.env,
         )
         dataset_properties = DatasetPropertiesClass(
             name=table_name,

datahub/ingestion/source/kafka_connect/common.py

@@ -141,12 +141,7 @@ def get_dataset_name(
     database_name: Optional[str],
     source_table: str,
 ) -> str:
-    if database_name:
-        dataset_name = database_name + "." + source_table
-    else:
-        dataset_name = source_table
-
-    return dataset_name
+    return database_name + "." + source_table if database_name else source_table
 
 
 def get_platform_instance(

datahub/ingestion/source/mlflow.py

@@ -1,17 +1,20 @@
+import time
 from dataclasses import dataclass
 from typing import Any, Callable, Iterable, List, Optional, TypeVar, Union
 
 from mlflow import MlflowClient
-from mlflow.entities import Run
+from mlflow.entities import Experiment, Run
 from mlflow.entities.model_registry import ModelVersion, RegisteredModel
 from mlflow.store.entities import PagedList
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.source_common import (
-    EnvConfigMixin,
+from datahub.api.entities.dataprocess.dataprocess_instance import (
+    DataProcessInstance,
 )
+from datahub.configuration.source_common import EnvConfigMixin
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -26,6 +29,7 @@ from datahub.ingestion.api.source import (
     SourceReport,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
@@ -35,20 +39,45 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.schema_classes import (
+    AuditStampClass,
+    ContainerClass,
+    DataPlatformInstanceClass,
+    DataProcessInstanceOutputClass,
+    DataProcessInstancePropertiesClass,
+    DataProcessInstanceRunEventClass,
+    DataProcessInstanceRunResultClass,
+    DataProcessRunStatusClass,
+    EdgeClass,
     GlobalTagsClass,
+    MetadataAttributionClass,
     MLHyperParamClass,
     MLMetricClass,
     MLModelGroupPropertiesClass,
     MLModelPropertiesClass,
+    MLTrainingRunPropertiesClass,
+    PlatformResourceInfoClass,
+    SubTypesClass,
     TagAssociationClass,
     TagPropertiesClass,
+    TimeStampClass,
+    VersionPropertiesClass,
     VersionTagClass,
     _Aspect,
 )
+from datahub.metadata.urns import (
+    DataPlatformUrn,
+    MlModelUrn,
+    VersionSetUrn,
+)
+from datahub.sdk.container import Container
 
 T = TypeVar("T")
 
 
+class ContainerKeyWithId(ContainerKey):
+    id: str
+
+
 class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     tracking_uri: Optional[str] = Field(
         default=None,
@@ -141,6 +170,7 @@ class MLflowSource(StatefulIngestionSourceBase):
 
     def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         yield from self._get_tags_workunits()
+        yield from self._get_experiment_workunits()
         yield from self._get_ml_model_workunits()
 
     def _get_tags_workunits(self) -> Iterable[MetadataWorkUnit]:
@@ -174,22 +204,162 @@ class MLflowSource(StatefulIngestionSourceBase):
             aspect=aspect,
         ).as_workunit()
 
-    def _get_ml_model_workunits(self) -> Iterable[MetadataWorkUnit]:
-        """
-        Traverse each Registered Model in Model Registry and generate a corresponding workunit.
-        """
-        registered_models = self._get_mlflow_registered_models()
-        for registered_model in registered_models:
-            yield self._get_ml_group_workunit(registered_model)
-            model_versions = self._get_mlflow_model_versions(registered_model)
-            for model_version in model_versions:
-                run = self._get_mlflow_run(model_version)
-                yield self._get_ml_model_properties_workunit(
-                    registered_model=registered_model,
-                    model_version=model_version,
-                    run=run,
-                )
-                yield self._get_global_tags_workunit(model_version=model_version)
+    def _get_experiment_workunits(self) -> Iterable[MetadataWorkUnit]:
+        experiments = self._get_mlflow_experiments()
+        for experiment in experiments:
+            yield from self._get_experiment_container_workunit(experiment)
+
+            runs = self._get_mlflow_runs_from_experiment(experiment)
+            if runs:
+                for run in runs:
+                    yield from self._get_run_workunits(experiment, run)
+
+    def _get_experiment_custom_properties(self, experiment):
+        experiment_custom_props = getattr(experiment, "tags", {}) or {}
+        experiment_custom_props.pop("mlflow.note.content", None)
+        experiment_custom_props["artifacts_location"] = experiment.artifact_location
+        return experiment_custom_props
+
+    def _get_experiment_container_workunit(
+        self, experiment: Experiment
+    ) -> Iterable[MetadataWorkUnit]:
+        experiment_container = Container(
+            container_key=ContainerKeyWithId(
+                platform=str(DataPlatformUrn(platform_name=self.platform)),
+                id=experiment.name,
+            ),
+            subtype=MLAssetSubTypes.MLFLOW_EXPERIMENT,
+            display_name=experiment.name,
+            description=experiment.tags.get("mlflow.note.content"),
+            extra_properties=self._get_experiment_custom_properties(experiment),
+        )
+
+        yield from experiment_container.as_workunits()
+
+    def _get_run_metrics(self, run: Run) -> List[MLMetricClass]:
+        return [
+            MLMetricClass(name=k, value=str(v)) for k, v in run.data.metrics.items()
+        ]
+
+    def _get_run_params(self, run: Run) -> List[MLHyperParamClass]:
+        return [
+            MLHyperParamClass(name=k, value=str(v)) for k, v in run.data.params.items()
+        ]
+
+    def _convert_run_result_type(
+        self, status: str
+    ) -> DataProcessInstanceRunResultClass:
+        if status == "FINISHED":
+            return DataProcessInstanceRunResultClass(
+                type="SUCCESS", nativeResultType=self.platform
+            )
+        elif status == "FAILED":
+            return DataProcessInstanceRunResultClass(
+                type="FAILURE", nativeResultType=self.platform
+            )
+        else:
+            return DataProcessInstanceRunResultClass(
+                type="SKIPPED", nativeResultType=self.platform
+            )
+
+    def _get_run_workunits(
+        self, experiment: Experiment, run: Run
+    ) -> Iterable[MetadataWorkUnit]:
+        experiment_key = ContainerKeyWithId(
+            platform=str(DataPlatformUrn(self.platform)), id=experiment.name
+        )
+
+        data_process_instance = DataProcessInstance(
+            id=run.info.run_id,
+            orchestrator=self.platform,
+            template_urn=None,
+        )
+
+        created_time = run.info.start_time or int(time.time() * 1000)
+        user_id = run.info.user_id if run.info.user_id else "mlflow"
+        guid_dict_user = {"platform": self.platform, "user": user_id}
+        platform_user_urn = (
+            f"urn:li:platformResource:{builder.datahub_guid(guid_dict_user)}"
+        )
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=platform_user_urn,
+            aspect=PlatformResourceInfoClass(
+                resourceType="user",
+                primaryKey=user_id,
+            ),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=DataProcessInstancePropertiesClass(
+                name=run.info.run_name or run.info.run_id,
+                created=AuditStampClass(
+                    time=created_time,
+                    actor=platform_user_urn,
+                ),
+                externalUrl=self._make_external_url_from_run(experiment, run),
+                customProperties=getattr(run, "tags", {}) or {},
+            ),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=ContainerClass(container=experiment_key.as_urn()),
+        ).as_workunit()
+
+        model_versions = self.get_mlflow_model_versions_from_run(run.info.run_id)
+        if model_versions:
+            model_version_urn = self._make_ml_model_urn(model_versions[0])
+            yield MetadataChangeProposalWrapper(
+                entityUrn=str(data_process_instance.urn),
+                aspect=DataProcessInstanceOutputClass(
+                    outputs=[],
+                    outputEdges=[
+                        EdgeClass(destinationUrn=model_version_urn),
+                    ],
+                ),
+            ).as_workunit()
+
+        metrics = self._get_run_metrics(run)
+        hyperparams = self._get_run_params(run)
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=MLTrainingRunPropertiesClass(
+                hyperParams=hyperparams,
+                trainingMetrics=metrics,
+                outputUrls=[run.info.artifact_uri],
+                id=run.info.run_id,
+            ),
+        ).as_workunit()
+
+        if run.info.end_time:
+            duration_millis = run.info.end_time - run.info.start_time
+
+            yield MetadataChangeProposalWrapper(
+                entityUrn=str(data_process_instance.urn),
+                aspect=DataProcessInstanceRunEventClass(
+                    status=DataProcessRunStatusClass.COMPLETE,
+                    timestampMillis=run.info.end_time,
+                    result=DataProcessInstanceRunResultClass(
+                        type=self._convert_run_result_type(run.info.status).type,
+                        nativeResultType=self.platform,
+                    ),
+                    durationMillis=duration_millis,
+                ),
+            ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=DataPlatformInstanceClass(
+                platform=str(DataPlatformUrn(self.platform))
+            ),
+        ).as_workunit()
+
+        yield MetadataChangeProposalWrapper(
+            entityUrn=str(data_process_instance.urn),
+            aspect=SubTypesClass(typeNames=[MLAssetSubTypes.MLFLOW_TRAINING_RUN]),
+        ).as_workunit()
 
     def _get_mlflow_registered_models(self) -> Iterable[RegisteredModel]:
         """
@@ -202,6 +372,19 @@ class MLflowSource(StatefulIngestionSourceBase):
         )
         return registered_models
 
+    def _get_mlflow_experiments(self) -> Iterable[Experiment]:
+        experiments: Iterable[Experiment] = self._traverse_mlflow_search_func(
+            search_func=self.client.search_experiments,
+        )
+        return experiments
+
+    def _get_mlflow_runs_from_experiment(self, experiment: Experiment) -> Iterable[Run]:
+        runs: Iterable[Run] = self._traverse_mlflow_search_func(
+            search_func=self.client.search_runs,
+            experiment_ids=[experiment.experiment_id],
+        )
+        return runs
+
     @staticmethod
     def _traverse_mlflow_search_func(
         search_func: Callable[..., PagedList[T]],
@@ -218,6 +401,13 @@ class MLflowSource(StatefulIngestionSourceBase):
             if not next_page_token:
                 return
 
+    def _get_latest_version(self, registered_model: RegisteredModel) -> Optional[str]:
+        return (
+            str(registered_model.latest_versions[0].version)
+            if registered_model.latest_versions
+            else None
+        )
+
     def _get_ml_group_workunit(
         self,
         registered_model: RegisteredModel,
@@ -229,7 +419,20 @@ class MLflowSource(StatefulIngestionSourceBase):
         ml_model_group_properties = MLModelGroupPropertiesClass(
             customProperties=registered_model.tags,
             description=registered_model.description,
-            createdAt=registered_model.creation_timestamp,
+            created=TimeStampClass(
+                time=registered_model.creation_timestamp, actor=None
+            ),
+            lastModified=TimeStampClass(
+                time=registered_model.last_updated_timestamp,
+                actor=None,
+            ),
+            version=VersionTagClass(
+                versionTag=self._get_latest_version(registered_model),
+                metadataAttribution=MetadataAttributionClass(
+                    time=registered_model.last_updated_timestamp,
+                    actor="urn:li:corpuser:datahub",
+                ),
+            ),
         )
         wu = self._create_workunit(
             urn=ml_model_group_urn,
@@ -259,6 +462,16 @@ class MLflowSource(StatefulIngestionSourceBase):
         )
         return model_versions
 
+    def get_mlflow_model_versions_from_run(self, run_id):
+        filter_string = f"run_id = '{run_id}'"
+
+        model_versions: Iterable[ModelVersion] = self._traverse_mlflow_search_func(
+            search_func=self.client.search_model_versions,
+            filter_string=filter_string,
+        )
+
+        return list(model_versions)
+
     def _get_mlflow_run(self, model_version: ModelVersion) -> Union[None, Run]:
         """
         Get a Run associated with a Model Version. Some MVs may exist without Run.
@@ -269,6 +482,67 @@ class MLflowSource(StatefulIngestionSourceBase):
         else:
             return None
 
+    def _get_ml_model_workunits(self) -> Iterable[MetadataWorkUnit]:
+        """
+        Traverse each Registered Model in Model Registry and generate a corresponding workunit.
+        """
+        registered_models = self._get_mlflow_registered_models()
+        for registered_model in registered_models:
+            version_set_urn = self._get_version_set_urn(registered_model)
+            yield self._get_ml_group_workunit(registered_model)
+            model_versions = self._get_mlflow_model_versions(registered_model)
+            for model_version in model_versions:
+                run = self._get_mlflow_run(model_version)
+                yield self._get_ml_model_properties_workunit(
+                    registered_model=registered_model,
+                    model_version=model_version,
+                    run=run,
+                )
+                yield self._get_ml_model_version_properties_workunit(
+                    model_version=model_version,
+                    version_set_urn=version_set_urn,
+                )
+                yield self._get_global_tags_workunit(model_version=model_version)
+
+    def _get_version_set_urn(self, registered_model: RegisteredModel) -> VersionSetUrn:
+        guid_dict = {"platform": self.platform, "name": registered_model.name}
+        version_set_urn = VersionSetUrn(
+            id=builder.datahub_guid(guid_dict),
+            entity_type=MlModelUrn.ENTITY_TYPE,
+        )
+
+        return version_set_urn
+
+    def _get_ml_model_version_properties_workunit(
+        self,
+        model_version: ModelVersion,
+        version_set_urn: VersionSetUrn,
+    ) -> MetadataWorkUnit:
+        ml_model_urn = self._make_ml_model_urn(model_version)
+
+        # get mlmodel name from ml model urn
+        ml_model_version_properties = VersionPropertiesClass(
+            version=VersionTagClass(
+                versionTag=str(model_version.version),
+                metadataAttribution=MetadataAttributionClass(
+                    time=model_version.creation_timestamp,
+                    actor="urn:li:corpuser:datahub",
+                ),
+            ),
+            versionSet=str(version_set_urn),
+            sortId=str(model_version.version).zfill(10),
+            aliases=[
+                VersionTagClass(versionTag=alias) for alias in model_version.aliases
+            ],
+        )
+
+        wu = MetadataChangeProposalWrapper(
+            entityUrn=str(ml_model_urn),
+            aspect=ml_model_version_properties,
+        ).as_workunit()
+
+        return wu
+
     def _get_ml_model_properties_workunit(
         self,
         registered_model: RegisteredModel,
@@ -282,28 +556,47 @@ class MLflowSource(StatefulIngestionSourceBase):
         """
         ml_model_group_urn = self._make_ml_model_group_urn(registered_model)
         ml_model_urn = self._make_ml_model_urn(model_version)
+
         if run:
-            hyperparams = [
-                MLHyperParamClass(name=k, value=str(v))
-                for k, v in run.data.params.items()
-            ]
-            training_metrics = [
-                MLMetricClass(name=k, value=str(v)) for k, v in run.data.metrics.items()
-            ]
+            # Use the same metrics and hyperparams from the run
+            hyperparams = self._get_run_params(run)
+            training_metrics = self._get_run_metrics(run)
+            run_urn = DataProcessInstance(
+                id=run.info.run_id,
+                orchestrator=self.platform,
+            ).urn
+
+            training_jobs = [str(run_urn)] if run_urn else []
         else:
             hyperparams = None
             training_metrics = None
+            training_jobs = []
+
+        created_time = model_version.creation_timestamp
+        created_actor = (
+            f"urn:li:platformResource:{model_version.user_id}"
+            if model_version.user_id
+            else None
+        )
+        model_version_tags = [f"{k}:{v}" for k, v in model_version.tags.items()]
+
         ml_model_properties = MLModelPropertiesClass(
             customProperties=model_version.tags,
             externalUrl=self._make_external_url(model_version),
+            lastModified=TimeStampClass(
+                time=model_version.last_updated_timestamp,
+                actor=None,
+            ),
             description=model_version.description,
-            date=model_version.creation_timestamp,
-            version=VersionTagClass(versionTag=str(model_version.version)),
+            created=TimeStampClass(
+                time=created_time,
+                actor=created_actor,
+            ),
             hyperParams=hyperparams,
             trainingMetrics=training_metrics,
-            # mlflow tags are dicts, but datahub tags are lists. currently use only keys from mlflow tags
-            tags=list(model_version.tags.keys()),
+            tags=model_version_tags,
             groups=[ml_model_group_urn],
+            trainingJobs=training_jobs,
         )
         wu = self._create_workunit(urn=ml_model_urn, aspect=ml_model_properties)
         return wu
@@ -337,6 +630,15 @@ class MLflowSource(StatefulIngestionSourceBase):
         else:
             return None
 
+    def _make_external_url_from_run(
+        self, experiment: Experiment, run: Run
+    ) -> Union[None, str]:
+        base_uri = self.client.tracking_uri
+        if base_uri.startswith("http"):
+            return f"{base_uri.rstrip('/')}/#/experiments/{experiment.experiment_id}/runs/{run.info.run_id}"
+        else:
+            return None
+
     def _get_global_tags_workunit(
         self,
         model_version: ModelVersion,
@@ -356,3 +658,8 @@ class MLflowSource(StatefulIngestionSourceBase):
             aspect=global_tags,
         )
         return wu
+
+    @classmethod
+    def create(cls, config_dict: dict, ctx: PipelineContext) -> "MLflowSource":
+        config = MLflowConfig.parse_obj(config_dict)
+        return cls(ctx, config)

datahub/ingestion/source/redshift/lineage.py

@@ -814,8 +814,8 @@ class RedshiftLineageExtractor:
 
         tablename = table.name
         if (
-            table.is_external_table
-            and schema.is_external_schema
+            table.is_external_table()
+            and schema.is_external_schema()
             and schema.external_platform
         ):
             # external_db_params = schema.option

datahub/ingestion/source/redshift/lineage_v2.py

@@ -403,8 +403,8 @@ class RedshiftSqlLineageV2(Closeable):
             for table in tables:
                 schema = db_schemas[self.database][schema_name]
                 if (
-                    table.is_external_table
-                    and schema.is_external_schema
+                    table.is_external_table()
+                    and schema.is_external_schema()
                     and schema.external_platform
                 ):
                     # external_db_params = schema.option
@@ -416,14 +416,26 @@ class RedshiftSqlLineageV2(Closeable):
                         platform_instance=self.config.platform_instance,
                         env=self.config.env,
                     )
-                    upstream_urn = mce_builder.make_dataset_urn_with_platform_instance(
-                        upstream_platform,
-                        f"{schema.external_database}.{table.name}",
-                        platform_instance=(
+                    if upstream_platform == self.platform:
+                        upstream_schema = schema.get_upstream_schema_name() or "public"
+                        upstream_dataset_name = (
+                            f"{schema.external_database}.{upstream_schema}.{table.name}"
+                        )
+                        upstream_platform_instance = self.config.platform_instance
+                    else:
+                        upstream_dataset_name = (
+                            f"{schema.external_database}.{table.name}"
+                        )
+                        upstream_platform_instance = (
                             self.config.platform_instance_map.get(upstream_platform)
                             if self.config.platform_instance_map
                             else None
-                        ),
+                        )
+
+                    upstream_urn = mce_builder.make_dataset_urn_with_platform_instance(
+                        upstream_platform,
+                        upstream_dataset_name,
+                        platform_instance=upstream_platform_instance,
                         env=self.config.env,
                     )