acryl-datahub 1.0.0.2rc4__py3-none-any.whl → 1.0.0.3__py3-none-any.whl

datahub/sdk/lineage_client.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import difflib
+import logging
+from typing import TYPE_CHECKING, List, Literal, Optional, Set, Union
+
+import datahub.metadata.schema_classes as models
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.errors import SdkUsageError
+from datahub.metadata.schema_classes import SchemaMetadataClass
+from datahub.metadata.urns import DatasetUrn, QueryUrn
+from datahub.sdk._shared import DatasetUrnOrStr
+from datahub.sdk._utils import DEFAULT_ACTOR_URN
+from datahub.sdk.dataset import ColumnLineageMapping, parse_cll_mapping
+from datahub.specific.dataset import DatasetPatchBuilder
+from datahub.sql_parsing.fingerprint_utils import generate_hash
+from datahub.utilities.ordered_set import OrderedSet
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+logger = logging.getLogger(__name__)
+
+_empty_audit_stamp = models.AuditStampClass(
+    time=0,
+    actor=DEFAULT_ACTOR_URN,
+)
+
+
+class LineageClient:
+    def __init__(self, client: DataHubClient):
+        self._client = client
+
+    def _get_fields_from_dataset_urn(self, dataset_urn: DatasetUrn) -> Set[str]:
+        schema_metadata = self._client._graph.get_aspect(
+            str(dataset_urn), SchemaMetadataClass
+        )
+        if schema_metadata is None:
+            return Set()
+
+        return {field.fieldPath for field in schema_metadata.fields}
+
+    @classmethod
+    def _get_strict_column_lineage(
+        cls,
+        upstream_fields: Set[str],
+        downstream_fields: Set[str],
+    ) -> ColumnLineageMapping:
+        """Find matches between upstream and downstream fields with case-insensitive matching."""
+        strict_column_lineage: ColumnLineageMapping = {}
+
+        # Create case-insensitive mapping of upstream fields
+        case_insensitive_map = {field.lower(): field for field in upstream_fields}
+
+        # Match downstream fields using case-insensitive comparison
+        for downstream_field in downstream_fields:
+            lower_field = downstream_field.lower()
+            if lower_field in case_insensitive_map:
+                # Use the original case of the upstream field
+                strict_column_lineage[downstream_field] = [
+                    case_insensitive_map[lower_field]
+                ]
+
+        return strict_column_lineage
+
+    @classmethod
+    def _get_fuzzy_column_lineage(
+        cls,
+        upstream_fields: Set[str],
+        downstream_fields: Set[str],
+    ) -> ColumnLineageMapping:
+        """Generate fuzzy matches between upstream and downstream fields."""
+
+        # Simple normalization function for better matching
+        def normalize(s: str) -> str:
+            return s.lower().replace("_", "")
+
+        # Create normalized lookup for upstream fields
+        normalized_upstream = {normalize(field): field for field in upstream_fields}
+
+        fuzzy_column_lineage = {}
+        for downstream_field in downstream_fields:
+            # Try exact match first
+            if downstream_field in upstream_fields:
+                fuzzy_column_lineage[downstream_field] = [downstream_field]
+                continue
+
+            # Try normalized match
+            norm_downstream = normalize(downstream_field)
+            if norm_downstream in normalized_upstream:
+                fuzzy_column_lineage[downstream_field] = [
+                    normalized_upstream[norm_downstream]
+                ]
+                continue
+
+            # If no direct match, find closest match using similarity
+            matches = difflib.get_close_matches(
+                norm_downstream,
+                normalized_upstream.keys(),
+                n=1,  # Return only the best match
+                cutoff=0.8,  # Adjust cutoff for sensitivity
+            )
+
+            if matches:
+                fuzzy_column_lineage[downstream_field] = [
+                    normalized_upstream[matches[0]]
+                ]
+
+        return fuzzy_column_lineage
+
+    def add_dataset_copy_lineage(
+        self,
+        *,
+        upstream: DatasetUrnOrStr,
+        downstream: DatasetUrnOrStr,
+        column_lineage: Union[
+            None, ColumnLineageMapping, Literal["auto_fuzzy", "auto_strict"]
+        ] = "auto_fuzzy",
+    ) -> None:
+        upstream = DatasetUrn.from_string(upstream)
+        downstream = DatasetUrn.from_string(downstream)
+
+        if column_lineage is None:
+            cll = None
+        elif column_lineage in ["auto_fuzzy", "auto_strict"]:
+            upstream_schema = self._get_fields_from_dataset_urn(upstream)
+            downstream_schema = self._get_fields_from_dataset_urn(downstream)
+            if column_lineage == "auto_fuzzy":
+                mapping = self._get_fuzzy_column_lineage(
+                    upstream_schema, downstream_schema
+                )
+            else:
+                mapping = self._get_strict_column_lineage(
+                    upstream_schema, downstream_schema
+                )
+            cll = parse_cll_mapping(
+                upstream=upstream,
+                downstream=downstream,
+                cll_mapping=mapping,
+            )
+        elif isinstance(column_lineage, dict):
+            cll = parse_cll_mapping(
+                upstream=upstream,
+                downstream=downstream,
+                cll_mapping=column_lineage,
+            )
+
+        updater = DatasetPatchBuilder(str(downstream))
+        updater.add_upstream_lineage(
+            models.UpstreamClass(
+                dataset=str(upstream),
+                type=models.DatasetLineageTypeClass.COPY,
+            )
+        )
+        for cl in cll or []:
+            updater.add_fine_grained_upstream_lineage(cl)
+
+        self._client.entities.update(updater)
+
+    def add_dataset_transform_lineage(
+        self,
+        *,
+        upstream: DatasetUrnOrStr,
+        downstream: DatasetUrnOrStr,
+        column_lineage: Optional[ColumnLineageMapping] = None,
+        query_text: Optional[str] = None,
+    ) -> None:
+        upstream = DatasetUrn.from_string(upstream)
+        downstream = DatasetUrn.from_string(downstream)
+
+        cll = None
+        if column_lineage is not None:
+            cll = parse_cll_mapping(
+                upstream=upstream,
+                downstream=downstream,
+                cll_mapping=column_lineage,
+            )
+
+        fields_involved = OrderedSet([str(upstream), str(downstream)])
+        if cll is not None:
+            for c in cll:
+                for field in c.upstreams or []:
+                    fields_involved.add(field)
+                for field in c.downstreams or []:
+                    fields_involved.add(field)
+
+        query_urn = None
+        query_entity = None
+        if query_text:
+            # Eventually we might want to use our regex-based fingerprinting instead.
+            fingerprint = generate_hash(query_text)
+            query_urn = QueryUrn(fingerprint).urn()
+
+            from datahub.sql_parsing.sql_parsing_aggregator import make_query_subjects
+
+            query_entity = MetadataChangeProposalWrapper.construct_many(
+                query_urn,
+                aspects=[
+                    models.QueryPropertiesClass(
+                        statement=models.QueryStatementClass(
+                            value=query_text, language=models.QueryLanguageClass.SQL
+                        ),
+                        source=models.QuerySourceClass.SYSTEM,
+                        created=_empty_audit_stamp,
+                        lastModified=_empty_audit_stamp,
+                    ),
+                    make_query_subjects(list(fields_involved)),
+                ],
+            )
+
+        updater = DatasetPatchBuilder(str(downstream))
+        updater.add_upstream_lineage(
+            models.UpstreamClass(
+                dataset=str(upstream),
+                type=models.DatasetLineageTypeClass.TRANSFORMED,
+                query=query_urn,
+            )
+        )
+        for cl in cll or []:
+            cl.query = query_urn
+            updater.add_fine_grained_upstream_lineage(cl)
+
+        # Throw if the dataset does not exist.
+        # We need to manually call .build() instead of reusing client.update()
+        # so that we make just one emit_mcps call.
+        if not self._client._graph.exists(updater.urn):
+            raise SdkUsageError(
+                f"Dataset {updater.urn} does not exist, and hence cannot be updated."
+            )
+        mcps: List[
+            Union[MetadataChangeProposalWrapper, models.MetadataChangeProposalClass]
+        ] = list(updater.build())
+        if query_entity:
+            mcps.extend(query_entity)
+        self._client._graph.emit_mcps(mcps)

datahub/sdk/main_client.py

@@ -4,8 +4,9 @@ from typing import Optional, overload
 
 from datahub.errors import SdkUsageError
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
-from datahub.ingestion.graph.config import DatahubClientConfig
+from datahub.ingestion.graph.config import ClientMode, DatahubClientConfig
 from datahub.sdk.entity_client import EntityClient
+from datahub.sdk.lineage_client import LineageClient
 from datahub.sdk.resolver_client import ResolverClient
 from datahub.sdk.search_client import SearchClient
 
@@ -83,7 +84,7 @@ class DataHubClient:
         # Inspired by the DockerClient.from_env() method.
         # TODO: This one also reads from ~/.datahubenv, so the "from_env" name might be a bit confusing.
         # That file is part of the "environment", but is not a traditional "env variable".
-        graph = get_default_graph()
+        graph = get_default_graph(ClientMode.SDK)
 
         return cls(graph=graph)
 
@@ -99,4 +100,6 @@ class DataHubClient:
     def search(self) -> SearchClient:
         return SearchClient(self)
 
-    # TODO: lineage client
+    @property
+    def lineage(self) -> LineageClient:
+        return LineageClient(self)

datahub/sdk/mlmodel.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, List, Optional, Sequence, Type, Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mce_builder import DEFAULT_ENV
+from datahub.metadata.schema_classes import (
+    AspectBag,
+    MLHyperParamClass,
+    MLMetricClass,
+    MLModelPropertiesClass,
+)
+from datahub.metadata.urns import (
+    DataProcessInstanceUrn,
+    MlModelGroupUrn,
+    MlModelUrn,
+    Urn,
+)
+from datahub.sdk._shared import (
+    DomainInputType,
+    HasDomain,
+    HasInstitutionalMemory,
+    HasOwnership,
+    HasPlatformInstance,
+    HasTags,
+    HasTerms,
+    HasVersion,
+    HyperParamsInputType,
+    LinksInputType,
+    MLTrainingJobInputType,
+    OwnersInputType,
+    TagsInputType,
+    TermsInputType,
+    TrainingMetricsInputType,
+    convert_hyper_params,
+    convert_training_metrics,
+    make_time_stamp,
+    parse_time_stamp,
+)
+from datahub.sdk.entity import Entity, ExtraAspectsType
+
+
+class MLModel(
+    HasPlatformInstance,
+    HasOwnership,
+    HasInstitutionalMemory,
+    HasTags,
+    HasTerms,
+    HasDomain,
+    HasVersion,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[MlModelUrn]:
+        return MlModelUrn
+
+    def __init__(
+        self,
+        id: str,
+        platform: str,
+        version: Optional[str] = None,
+        aliases: Optional[List[str]] = None,
+        platform_instance: Optional[str] = None,
+        env: str = DEFAULT_ENV,
+        name: Optional[str] = None,
+        description: Optional[str] = None,
+        training_metrics: Optional[TrainingMetricsInputType] = None,
+        hyper_params: Optional[HyperParamsInputType] = None,
+        external_url: Optional[str] = None,
+        custom_properties: Optional[Dict[str, str]] = None,
+        created: Optional[datetime] = None,
+        last_modified: Optional[datetime] = None,
+        owners: Optional[OwnersInputType] = None,
+        links: Optional[LinksInputType] = None,
+        tags: Optional[TagsInputType] = None,
+        terms: Optional[TermsInputType] = None,
+        domain: Optional[DomainInputType] = None,
+        model_group: Optional[Union[str, MlModelGroupUrn]] = None,
+        training_jobs: Optional[MLTrainingJobInputType] = None,
+        downstream_jobs: Optional[MLTrainingJobInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
+    ):
+        urn = MlModelUrn(platform=platform, name=id, env=env)
+        super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
+
+        self._set_platform_instance(urn.platform, platform_instance)
+
+        self._ensure_model_props()
+
+        if version is not None:
+            self.set_version(version)
+        if name is not None:
+            self.set_name(name)
+        if aliases is not None:
+            self.set_version_aliases(aliases)
+        if description is not None:
+            self.set_description(description)
+        if training_metrics is not None:
+            self.set_training_metrics(training_metrics)
+        if hyper_params is not None:
+            self.set_hyper_params(hyper_params)
+        if external_url is not None:
+            self.set_external_url(external_url)
+        if custom_properties is not None:
+            self.set_custom_properties(custom_properties)
+        if created is not None:
+            self.set_created(created)
+        if last_modified is not None:
+            self.set_last_modified(last_modified)
+
+        if owners is not None:
+            self.set_owners(owners)
+        if links is not None:
+            self.set_links(links)
+        if tags is not None:
+            self.set_tags(tags)
+        if terms is not None:
+            self.set_terms(terms)
+        if domain is not None:
+            self.set_domain(domain)
+        if model_group is not None:
+            self.set_model_group(model_group)
+        if training_jobs is not None:
+            self.set_training_jobs(training_jobs)
+        if downstream_jobs is not None:
+            self.set_downstream_jobs(downstream_jobs)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: AspectBag) -> Self:
+        assert isinstance(urn, MlModelUrn)
+        entity = cls(
+            id=urn.name,
+            platform=urn.platform,
+            env=urn.env,
+        )
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> MlModelUrn:
+        return self._urn  # type: ignore
+
+    def _ensure_model_props(
+        self,
+    ) -> MLModelPropertiesClass:
+        return self._setdefault_aspect(MLModelPropertiesClass())
+
+    @property
+    def name(self) -> Optional[str]:
+        return self._ensure_model_props().name
+
+    def set_name(self, name: str) -> None:
+        self._ensure_model_props().name = name
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_model_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_model_props().description = description
+
+    @property
+    def external_url(self) -> Optional[str]:
+        return self._ensure_model_props().externalUrl
+
+    def set_external_url(self, external_url: str) -> None:
+        self._ensure_model_props().externalUrl = external_url
+
+    @property
+    def custom_properties(self) -> Optional[Dict[str, str]]:
+        return self._ensure_model_props().customProperties
+
+    def set_custom_properties(self, custom_properties: Dict[str, str]) -> None:
+        self._ensure_model_props().customProperties = custom_properties
+
+    @property
+    def created(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_props().created)
+
+    def set_created(self, created: datetime) -> None:
+        self._ensure_model_props().created = make_time_stamp(created)
+
+    @property
+    def last_modified(self) -> Optional[datetime]:
+        return parse_time_stamp(self._ensure_model_props().lastModified)
+
+    def set_last_modified(self, last_modified: datetime) -> None:
+        self._ensure_model_props().lastModified = make_time_stamp(last_modified)
+
+    @property
+    def training_metrics(self) -> Optional[List[MLMetricClass]]:
+        return self._ensure_model_props().trainingMetrics
+
+    def set_training_metrics(self, metrics: TrainingMetricsInputType) -> None:
+        self._ensure_model_props().trainingMetrics = convert_training_metrics(metrics)
+
+    def add_training_metrics(self, metrics: TrainingMetricsInputType) -> None:
+        props = self._ensure_model_props()
+        if props.trainingMetrics is None:
+            props.trainingMetrics = []
+        if isinstance(metrics, list):
+            props.trainingMetrics.extend(
+                [
+                    MLMetricClass(name=metric.name, value=metric.value)
+                    for metric in metrics
+                ]
+            )
+        else:
+            # For dictionary case, use the key as name and value as value
+            for name, value in metrics.items():
+                props.trainingMetrics.append(MLMetricClass(name=name, value=value))
+
+    @property
+    def hyper_params(self) -> Optional[List[MLHyperParamClass]]:
+        return self._ensure_model_props().hyperParams
+
+    def set_hyper_params(self, params: HyperParamsInputType) -> None:
+        self._ensure_model_props().hyperParams = convert_hyper_params(params)
+
+    def add_hyper_params(self, params: HyperParamsInputType) -> None:
+        props = self._ensure_model_props()
+        if props.hyperParams is None:
+            props.hyperParams = []
+        if isinstance(params, list):
+            props.hyperParams.extend(
+                [
+                    MLHyperParamClass(name=param.name, value=param.value)
+                    for param in params
+                ]
+            )
+        else:
+            # For dictionary case, iterate through key-value pairs
+            for name, value in params.items():
+                props.hyperParams.append(MLHyperParamClass(name=name, value=value))
+
+    @property
+    def model_group(self) -> Optional[str]:
+        props = self._ensure_model_props()
+        groups = props.groups
+        if groups is None or len(groups) == 0:
+            return None
+        return groups[0]
+
+    def set_model_group(self, group: Union[str, MlModelGroupUrn]) -> None:
+        self._ensure_model_props().groups = [str(group)]
+
+    @property
+    def training_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_props().trainingJobs
+
+    def set_training_jobs(self, training_jobs: MLTrainingJobInputType) -> None:
+        self._ensure_model_props().trainingJobs = [str(job) for job in training_jobs]
+
+    def add_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.trainingJobs is None:
+            props.trainingJobs = []
+        props.trainingJobs.append(str(training_job))
+
+    def remove_training_job(
+        self, training_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.trainingJobs is not None:
+            job_str = str(training_job)
+            props.trainingJobs = [job for job in props.trainingJobs if job != job_str]
+
+    @property
+    def downstream_jobs(self) -> Optional[List[str]]:
+        return self._ensure_model_props().downstreamJobs
+
+    def set_downstream_jobs(
+        self, downstream_jobs: Sequence[Union[str, DataProcessInstanceUrn]]
+    ) -> None:
+        self._ensure_model_props().downstreamJobs = [
+            str(job) for job in downstream_jobs
+        ]
+
+    def add_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.downstreamJobs is None:
+            props.downstreamJobs = []
+        props.downstreamJobs.append(str(downstream_job))
+
+    def remove_downstream_job(
+        self, downstream_job: Union[str, DataProcessInstanceUrn]
+    ) -> None:
+        props = self._ensure_model_props()
+        if props.downstreamJobs is not None:
+            job_str = str(downstream_job)
+            props.downstreamJobs = [
+                job for job in props.downstreamJobs if job != job_str
+            ]