acryl-datahub 1.0.0.1rc1__py3-none-any.whl → 1.0.0.1rc2__py3-none-any.whl

datahub/ingestion/source/mlflow.py

@@ -1,9 +1,10 @@
+import json
 import time
 from dataclasses import dataclass
-from typing import Any, Callable, Iterable, List, Optional, TypeVar, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, TypeVar, Union
 
 from mlflow import MlflowClient
-from mlflow.entities import Experiment, Run
+from mlflow.entities import Dataset as MlflowDataset, Experiment, Run
 from mlflow.entities.model_registry import ModelVersion, RegisteredModel
 from mlflow.store.entities import PagedList
 from pydantic.fields import Field
@@ -29,6 +30,7 @@ from datahub.ingestion.api.source import (
     SourceReport,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.data_platforms import KNOWN_VALID_PLATFORM_NAMES
 from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -42,6 +44,7 @@ from datahub.metadata.schema_classes import (
     AuditStampClass,
     ContainerClass,
     DataPlatformInstanceClass,
+    DataProcessInstanceInputClass,
     DataProcessInstanceOutputClass,
     DataProcessInstancePropertiesClass,
     DataProcessInstanceRunEventClass,
@@ -60,16 +63,15 @@ from datahub.metadata.schema_classes import (
     TagAssociationClass,
     TagPropertiesClass,
     TimeStampClass,
+    UpstreamClass,
+    UpstreamLineageClass,
     VersionPropertiesClass,
     VersionTagClass,
     _Aspect,
 )
-from datahub.metadata.urns import (
-    DataPlatformUrn,
-    MlModelUrn,
-    VersionSetUrn,
-)
+from datahub.metadata.urns import DataPlatformUrn, DatasetUrn, MlModelUrn, VersionSetUrn
 from datahub.sdk.container import Container
+from datahub.sdk.dataset import Dataset
 
 T = TypeVar("T")
 
@@ -105,6 +107,13 @@ class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
             " If neither is set, external URLs are not generated."
         ),
     )
+    materialize_dataset_inputs: Optional[bool] = Field(
+        default=False,
+        description="Whether to materialize dataset inputs for each run",
+    )
+    source_mapping_to_platform: Optional[dict] = Field(
+        default=None, description="Mapping of source type to datahub platform"
+    )
 
 
 @dataclass
@@ -213,6 +222,7 @@ class MLflowSource(StatefulIngestionSourceBase):
             if runs:
                 for run in runs:
                     yield from self._get_run_workunits(experiment, run)
+                    yield from self._get_dataset_input_workunits(run)
 
     def _get_experiment_custom_properties(self, experiment):
         experiment_custom_props = getattr(experiment, "tags", {}) or {}
@@ -262,6 +272,183 @@ class MLflowSource(StatefulIngestionSourceBase):
                 type="SKIPPED", nativeResultType=self.platform
             )
 
+    def _get_dataset_schema(
+        self, dataset: MlflowDataset
+    ) -> Optional[List[Tuple[str, str]]]:
+        try:
+            schema_dict = json.loads(dataset.schema)
+        except json.JSONDecodeError:
+            self.report.warning(
+                title="Failed to load dataset schema",
+                message="Schema metadata will be missing due to a JSON parsing error.",
+                context=f"Dataset: {dataset.name}, Schema: {dataset.schema}",
+            )
+            return None
+
+        if "mlflow_colspec" in schema_dict:
+            try:
+                return [
+                    (field["name"], field["type"])
+                    for field in schema_dict["mlflow_colspec"]
+                ]
+            except (KeyError, TypeError):
+                return None
+        # If the schema is not formatted, return None
+        return None
+
+    def _get_external_dataset_urn(self, platform: str, dataset_name: str) -> str:
+        """
+        Get the URN for an external dataset.
+        Args:
+            platform: The platform of the external dataset (e.g., 's3', 'bigquery')
+            dataset: The MLflow dataset
+        Returns:
+            str: The URN of the external dataset
+        """
+        return str(DatasetUrn(platform=platform, name=dataset_name))
+
+    def _get_dataset_input_workunits(self, run: Run) -> Iterable[MetadataWorkUnit]:
+        """
+        Generate workunits for dataset inputs in a run.
+
+        For each dataset input:
+        1. If source type is 'local' or 'code':
+           - Create a local dataset reference
+        2. Otherwise:
+           - If materialization is enabled:
+             - Create a hosted dataset and a dataset reference with upstream
+           - If materialization is not enabled:
+             - Create a dataset reference and add upstream if dataset exists
+        3. Add all dataset references as upstreams for the run
+        """
+        run_urn = DataProcessInstance(
+            id=run.info.run_id,
+            orchestrator=self.platform,
+        ).urn
+
+        dataset_reference_urns = []
+
+        for dataset_input in run.inputs.dataset_inputs:
+            dataset = dataset_input.dataset
+            source_type = dataset.source_type
+            dataset_tags = {k[1]: v[1] for k, v in dataset_input.tags}
+
+            # Prepare dataset properties
+            custom_properties = dataset_tags
+            formatted_schema = self._get_dataset_schema(dataset)
+            if formatted_schema is None:
+                custom_properties["schema"] = dataset.schema
+
+            # Handle local/code datasets
+            if source_type in ("local", "code"):
+                local_dataset = Dataset(
+                    platform=self.platform,
+                    name=dataset.name,
+                    schema=formatted_schema,
+                    custom_properties=custom_properties,
+                )
+                yield from local_dataset.as_workunits()
+                dataset_reference_urns.append(local_dataset.urn)
+                continue
+
+            # Handle hosted datasets
+            formatted_platform = self._get_dataset_platform_from_source_type(
+                source_type
+            )
+
+            # Validate platform if materialization is enabled
+            if self.config.materialize_dataset_inputs:
+                if not formatted_platform:
+                    self.report.failure(
+                        title="Unable to materialize dataset inputs",
+                        message=f"No mapping dataPlatform found for dataset input source type '{source_type}'",
+                        context=f"please add `materialize_dataset_inputs.source_mapping_to_platform` in config "
+                        f"(e.g. '{source_type}': 'snowflake')",
+                    )
+                    continue
+                # Create hosted dataset
+                hosted_dataset = Dataset(
+                    platform=formatted_platform,
+                    name=dataset.name,
+                    schema=formatted_schema,
+                    custom_properties=dataset_tags,
+                )
+                yield from hosted_dataset.as_workunits()
+
+            # Create dataset reference with upstream
+            hosted_dataset_reference = Dataset(
+                platform=self.platform,
+                name=dataset.name,
+                schema=formatted_schema,
+                custom_properties=dataset_tags,
+                upstreams=UpstreamLineageClass(
+                    upstreams=[
+                        UpstreamClass(
+                            self._get_external_dataset_urn(
+                                formatted_platform, dataset.name
+                            ),
+                            type="COPY",
+                        )
+                    ]
+                )
+                if formatted_platform
+                else None,
+            )
+            dataset_reference_urns.append(hosted_dataset_reference.urn)
+            yield from hosted_dataset_reference.as_workunits()
+
+        # Add dataset references as upstreams for the run
+        if dataset_reference_urns:
+            input_edges = [
+                EdgeClass(destinationUrn=str(dataset_ref_urn))
+                for dataset_ref_urn in dataset_reference_urns
+            ]
+            yield MetadataChangeProposalWrapper(
+                entityUrn=str(run_urn),
+                aspect=DataProcessInstanceInputClass(inputs=[], inputEdges=input_edges),
+            ).as_workunit()
+
+    def _get_dataset_platform_from_source_type(self, source_type: str) -> Optional[str]:
+        """
+        Map MLflow source type to DataHub platform.
+
+        Priority:
+        1. User-provided mapping in config
+        2. Internal mapping
+        3. Direct platform match from list of supported platforms
+        """
+        source_type = source_type.lower()
+
+        # User-provided mapping
+        platform = self._get_platform_from_user_mapping(source_type)
+        if platform:
+            return platform
+
+        # Internal mapping
+        if source_type == "gs":
+            return "gcs"
+
+        # Check direct platform match
+        if self._is_valid_platform(source_type):
+            return source_type
+
+        return None
+
+    def _get_platform_from_user_mapping(self, source_type: str) -> Optional[str]:
+        """
+        Get platform from user-provided mapping in config.
+        Returns None if mapping is invalid or platform is not supported.
+        """
+        source_mapping = self.config.source_mapping_to_platform
+        if not source_mapping:
+            return None
+
+        platform = source_mapping.get(source_type)
+        if not platform:
+            return None
+
+        return platform
+
     def _get_run_workunits(
         self, experiment: Experiment, run: Run
     ) -> Iterable[MetadataWorkUnit]:
@@ -659,6 +846,10 @@ class MLflowSource(StatefulIngestionSourceBase):
         )
         return wu
 
+    def _is_valid_platform(self, platform: Optional[str]) -> bool:
+        """Check if platform is registered as a source plugin"""
+        return platform in KNOWN_VALID_PLATFORM_NAMES
+
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "MLflowSource":
         config = MLflowConfig.parse_obj(config_dict)

datahub/ingestion/source/mode.py

@@ -33,6 +33,7 @@ from datahub.emitter.mcp_builder import (
     add_dataset_to_container,
     gen_containers,
 )
+from datahub.emitter.request_helper import make_curl_command
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -339,7 +340,8 @@ class ModeSource(StatefulIngestionSourceBase):
 
         # Test the connection
         try:
-            self._get_request_json(f"{self.config.connect_uri}/api/verify")
+            key_info = self._get_request_json(f"{self.config.connect_uri}/api/verify")
+            logger.debug(f"Auth info: {key_info}")
         except ModeRequestError as e:
             self.report.report_failure(
                 title="Failed to Connect",
@@ -1485,12 +1487,17 @@ class ModeSource(StatefulIngestionSourceBase):
 
         @r.wraps
         def get_request():
+            curl_command = make_curl_command(self.session, "GET", url, "")
+            logger.debug(f"Issuing request; curl equivalent: {curl_command}")
+
             try:
                 response = self.session.get(
                     url, timeout=self.config.api_options.timeout
                 )
                 if response.status_code == 204:  # No content, don't parse json
                     return {}
+
+                response.raise_for_status()
                 return response.json()
             except HTTPError as http_error:
                 error_response = http_error.response
@@ -1501,6 +1508,9 @@ class ModeSource(StatefulIngestionSourceBase):
                         time.sleep(float(sleep_time))
                     raise HTTPError429 from None
 
+                logger.debug(
+                    f"Error response ({error_response.status_code}): {error_response.text}"
+                )
                 raise http_error
 
         return get_request()

datahub/ingestion/source/openapi.py

@@ -2,13 +2,14 @@ import logging
 import time
 import warnings
 from abc import ABC
-from typing import Dict, Iterable, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple
 
 from pydantic import validator
 from pydantic.fields import Field
 
 from datahub.configuration.common import ConfigModel
 from datahub.emitter.mce_builder import make_tag_urn
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SourceCapability,
@@ -20,6 +21,7 @@ from datahub.ingestion.api.decorators import (
 )
 from datahub.ingestion.api.source import Source, SourceReport
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import DatasetSubTypes
 from datahub.ingestion.source.openapi_parser import (
     clean_url,
     compose_url_attr,
@@ -32,14 +34,13 @@ from datahub.ingestion.source.openapi_parser import (
     set_metadata,
     try_guessing,
 )
-from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
-from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import (
     AuditStampClass,
     DatasetPropertiesClass,
     GlobalTagsClass,
     InstitutionalMemoryClass,
     InstitutionalMemoryMetadataClass,
+    SubTypesClass,
     TagAssociationClass,
 )
 
@@ -222,8 +223,9 @@ class APISource(Source, ABC):
 
     def init_dataset(
         self, endpoint_k: str, endpoint_dets: dict
-    ) -> Tuple[DatasetSnapshot, str]:
+    ) -> Tuple[str, str, List[MetadataWorkUnit]]:
         config = self.config
+        workunits = []
 
         dataset_name = endpoint_k[1:].replace("/", ".")
 
@@ -233,22 +235,27 @@ class APISource(Source, ABC):
         else:
             dataset_name = "root"
 
-        dataset_snapshot = DatasetSnapshot(
-            urn=f"urn:li:dataset:(urn:li:dataPlatform:{self.platform},{config.name}.{dataset_name},PROD)",
-            aspects=[],
-        )
+        dataset_urn = f"urn:li:dataset:(urn:li:dataPlatform:{self.platform},{config.name}.{dataset_name},PROD)"
 
-        # adding description
-        dataset_properties = DatasetPropertiesClass(
+        # Create dataset properties aspect
+        properties = DatasetPropertiesClass(
             description=endpoint_dets["description"], customProperties={}
         )
-        dataset_snapshot.aspects.append(dataset_properties)
+        wu = MetadataWorkUnit(
+            id=dataset_name,
+            mcp=MetadataChangeProposalWrapper(entityUrn=dataset_urn, aspect=properties),
+        )
+        workunits.append(wu)
 
-        # adding tags
+        # Create tags aspect
         tags_str = [make_tag_urn(t) for t in endpoint_dets["tags"]]
         tags_tac = [TagAssociationClass(t) for t in tags_str]
         gtc = GlobalTagsClass(tags_tac)
-        dataset_snapshot.aspects.append(gtc)
+        wu = MetadataWorkUnit(
+            id=f"{dataset_name}-tags",
+            mcp=MetadataChangeProposalWrapper(entityUrn=dataset_urn, aspect=gtc),
+        )
+        workunits.append(wu)
 
         # the link will appear in the "documentation"
         link_url = clean_url(config.url + self.url_basepath + endpoint_k)
@@ -260,17 +267,25 @@ class APISource(Source, ABC):
             url=link_url, description=link_description, createStamp=creation
         )
         inst_memory = InstitutionalMemoryClass([link_metadata])
-        dataset_snapshot.aspects.append(inst_memory)
+        wu = MetadataWorkUnit(
+            id=f"{dataset_name}-docs",
+            mcp=MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn, aspect=inst_memory
+            ),
+        )
+        workunits.append(wu)
 
-        return dataset_snapshot, dataset_name
+        # Create subtype aspect
+        sub_types = SubTypesClass(typeNames=[DatasetSubTypes.API_ENDPOINT])
+        wu = MetadataWorkUnit(
+            id=f"{dataset_name}-subtype",
+            mcp=MetadataChangeProposalWrapper(entityUrn=dataset_urn, aspect=sub_types),
+        )
+        workunits.append(wu)
 
-    def build_wu(
-        self, dataset_snapshot: DatasetSnapshot, dataset_name: str
-    ) -> ApiWorkUnit:
-        mce = MetadataChangeEvent(proposedSnapshot=dataset_snapshot)
-        return ApiWorkUnit(id=dataset_name, mce=mce)
+        return dataset_name, dataset_urn, workunits
 
-    def get_workunits_internal(self) -> Iterable[ApiWorkUnit]:
+    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
         config = self.config
 
         sw_dict = self.config.get_swagger()
@@ -294,16 +309,24 @@ class APISource(Source, ABC):
             if endpoint_k in config.ignore_endpoints:
                 continue
 
-            dataset_snapshot, dataset_name = self.init_dataset(
+            # Initialize dataset and get common aspects
+            dataset_name, dataset_urn, workunits = self.init_dataset(
                 endpoint_k, endpoint_dets
             )
+            for wu in workunits:
+                yield wu
 
-            # adding dataset fields
+            # Handle schema metadata if available
             if "data" in endpoint_dets.keys():
                 # we are lucky! data is defined in the swagger for this endpoint
                 schema_metadata = set_metadata(dataset_name, endpoint_dets["data"])
-                dataset_snapshot.aspects.append(schema_metadata)
-                yield self.build_wu(dataset_snapshot, dataset_name)
+                wu = MetadataWorkUnit(
+                    id=f"{dataset_name}-schema",
+                    mcp=MetadataChangeProposalWrapper(
+                        entityUrn=dataset_urn, aspect=schema_metadata
+                    ),
+                )
+                yield wu
             elif endpoint_dets["method"] != "get":
                 self.report.report_warning(
                     title="Failed to Extract Endpoint Metadata",
@@ -338,9 +361,13 @@ class APISource(Source, ABC):
                             context=f"Endpoint Type: {endpoint_k}, Name: {dataset_name}",
                         )
                     schema_metadata = set_metadata(dataset_name, fields2add)
-                    dataset_snapshot.aspects.append(schema_metadata)
-
-                    yield self.build_wu(dataset_snapshot, dataset_name)
+                    wu = MetadataWorkUnit(
+                        id=f"{dataset_name}-schema",
+                        mcp=MetadataChangeProposalWrapper(
+                            entityUrn=dataset_urn, aspect=schema_metadata
+                        ),
+                    )
+                    yield wu
                 else:
                     self.report_bad_responses(response.status_code, type=endpoint_k)
             else:
@@ -369,9 +396,13 @@ class APISource(Source, ABC):
                                 context=f"Endpoint Type: {endpoint_k}, Name: {dataset_name}",
                             )
                         schema_metadata = set_metadata(dataset_name, fields2add)
-                        dataset_snapshot.aspects.append(schema_metadata)
-
-                        yield self.build_wu(dataset_snapshot, dataset_name)
+                        wu = MetadataWorkUnit(
+                            id=f"{dataset_name}-schema",
+                            mcp=MetadataChangeProposalWrapper(
+                                entityUrn=dataset_urn, aspect=schema_metadata
+                            ),
+                        )
+                        yield wu
                     else:
                         self.report_bad_responses(response.status_code, type=endpoint_k)
                 else:
@@ -400,9 +431,13 @@ class APISource(Source, ABC):
                                 context=f"Endpoint Type: {endpoint_k}, Name: {dataset_name}",
                             )
                         schema_metadata = set_metadata(dataset_name, fields2add)
-                        dataset_snapshot.aspects.append(schema_metadata)
-
-                        yield self.build_wu(dataset_snapshot, dataset_name)
+                        wu = MetadataWorkUnit(
+                            id=f"{dataset_name}-schema",
+                            mcp=MetadataChangeProposalWrapper(
+                                entityUrn=dataset_urn, aspect=schema_metadata
+                            ),
+                        )
+                        yield wu
                     else:
                         self.report_bad_responses(response.status_code, type=endpoint_k)
 

datahub/ingestion/source/powerbi/powerbi.py

@@ -666,6 +666,7 @@ class Mapper:
         workspace: powerbi_data_classes.Workspace,
         chart_mcps: List[MetadataChangeProposalWrapper],
         user_mcps: List[MetadataChangeProposalWrapper],
+        dashboard_edges: List[EdgeClass],
     ) -> List[MetadataChangeProposalWrapper]:
         """
         Map PowerBi dashboard to Datahub dashboard
@@ -695,6 +696,7 @@ class Mapper:
             lastModified=ChangeAuditStamps(),
             dashboardUrl=dashboard.webUrl,
             customProperties={**chart_custom_properties(dashboard)},
+            dashboards=dashboard_edges,
         )
 
         info_mcp = self.new_mcp(
@@ -933,7 +935,7 @@ class Mapper:
         dashboard: powerbi_data_classes.Dashboard,
         workspace: powerbi_data_classes.Workspace,
     ) -> List[EquableMetadataWorkUnit]:
-        mcps = []
+        mcps: List[MetadataChangeProposalWrapper] = []
 
         logger.info(
             f"Converting dashboard={dashboard.displayName} to datahub dashboard"
@@ -945,9 +947,30 @@ class Mapper:
         )
         # Convert tiles to charts
         ds_mcps, chart_mcps = self.to_datahub_chart(dashboard.tiles, workspace)
+
+        # collect all downstream reports (dashboards)
+        dashboard_edges = []
+        for t in dashboard.tiles:
+            if t.report:
+                dashboard_urn = builder.make_dashboard_urn(
+                    platform=self.__config.platform_name,
+                    platform_instance=self.__config.platform_instance,
+                    name=t.report.get_urn_part(),
+                )
+                edge = EdgeClass(
+                    destinationUrn=dashboard_urn,
+                )
+                dashboard_edges.append(edge)
+
         # Lets convert dashboard to datahub dashboard
         dashboard_mcps: List[MetadataChangeProposalWrapper] = (
-            self.to_datahub_dashboard_mcp(dashboard, workspace, chart_mcps, user_mcps)
+            self.to_datahub_dashboard_mcp(
+                dashboard=dashboard,
+                workspace=workspace,
+                chart_mcps=chart_mcps,
+                user_mcps=user_mcps,
+                dashboard_edges=dashboard_edges,
+            )
         )
 
         # Now add MCPs in sequence
@@ -1054,7 +1077,6 @@ class Mapper:
         report: powerbi_data_classes.Report,
         chart_mcps: List[MetadataChangeProposalWrapper],
         user_mcps: List[MetadataChangeProposalWrapper],
-        dashboard_edges: List[EdgeClass],
     ) -> List[MetadataChangeProposalWrapper]:
         """
         Map PowerBi report to Datahub dashboard
@@ -1076,7 +1098,6 @@ class Mapper:
             charts=chart_urn_list,
             lastModified=ChangeAuditStamps(),
             dashboardUrl=report.webUrl,
-            dashboards=dashboard_edges,
         )
 
         info_mcp = self.new_mcp(
@@ -1170,27 +1191,12 @@ class Mapper:
         ds_mcps = self.to_datahub_dataset(report.dataset, workspace)
         chart_mcps = self.pages_to_chart(report.pages, workspace, ds_mcps)
 
-        # find all dashboards with a Tile referencing this report
-        downstream_dashboards_edges = []
-        for d in workspace.dashboards.values():
-            if any(t.report_id == report.id for t in d.tiles):
-                dashboard_urn = builder.make_dashboard_urn(
-                    platform=self.__config.platform_name,
-                    platform_instance=self.__config.platform_instance,
-                    name=d.get_urn_part(),
-                )
-                edge = EdgeClass(
-                    destinationUrn=dashboard_urn,
-                    sourceUrn=None,
-                    created=None,
-                    lastModified=None,
-                    properties=None,
-                )
-                downstream_dashboards_edges.append(edge)
-
         # Let's convert report to datahub dashboard
         report_mcps = self.report_to_dashboard(
-            workspace, report, chart_mcps, user_mcps, downstream_dashboards_edges
+            workspace=workspace,
+            report=report,
+            chart_mcps=chart_mcps,
+            user_mcps=user_mcps,
         )
 
         # Now add MCPs in sequence

datahub/ingestion/source/s3/source.py

@@ -945,6 +945,17 @@ class S3Source(StatefulIngestionSourceBase):
                     for f in list_folders(
                         bucket_name, f"{folder}", self.source_config.aws_config
                     ):
+                        table_path = self.create_s3_path(bucket_name, f)
+                        table_name, _ = path_spec.extract_table_name_and_path(
+                            table_path
+                        )
+                        if not path_spec.tables_filter_pattern.allowed(table_name):
+                            logger.debug(
+                                f"Table '{table_name}' not allowed and skipping"
+                            )
+                            self.report.report_file_dropped(table_path)
+                            continue
+
                         dirs_to_process = []
                         logger.info(f"Processing folder: {f}")
                         if path_spec.traversal_method == FolderTraversalMethod.ALL: