acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.2.0.1__py3-none-any.whl

datahub/sdk/main_client.py

@@ -7,16 +7,8 @@ from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
 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
 
-try:
-    from acryl_datahub_cloud._sdk_extras import (  # type: ignore[import-not-found]
-        AssertionsClient,
-    )
-except ImportError:
-    AssertionsClient = None
-
 
 class DataHubClient:
     """Main client for interacting with DataHub.
@@ -74,7 +66,12 @@ class DataHubClient:
         self._graph.test_connection()
 
     @classmethod
-    def from_env(cls) -> "DataHubClient":
+    def from_env(
+        cls,
+        *,
+        client_mode: ClientMode = ClientMode.SDK,
+        datahub_component: Optional[str] = None,
+    ) -> "DataHubClient":
         """Initialize a DataHubClient from the environment variables or ~/.datahubenv file.
 
         This will first check DATAHUB_GMS_URL and DATAHUB_GMS_TOKEN. If not present,
@@ -84,6 +81,10 @@ class DataHubClient:
         If you're looking to specify the server/token in code, use the
         DataHubClient(server=..., token=...) constructor instead.
 
+        Args:
+            client_mode: [internal] The client mode to use. Defaults to "SDK".
+            datahub_component: [internal] The DataHub component name to include in the user agent.
+
         Returns:
             A DataHubClient instance.
         """
@@ -91,7 +92,10 @@ 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(ClientMode.SDK)
+        graph = get_default_graph(
+            client_mode=client_mode,
+            datahub_component=datahub_component,
+        )
 
         return cls(graph=graph)
 
@@ -100,7 +104,15 @@ class DataHubClient:
         return EntityClient(self)
 
     @property
-    def resolve(self) -> ResolverClient:
+    def resolve(self):  # type: ignore[report-untyped-call]  # Not available due to circular import issues
+        try:
+            from acryl_datahub_cloud.sdk import (  # type: ignore[import-not-found]
+                ResolverClient,
+            )
+        except ImportError:
+            from datahub.sdk.resolver_client import (  # type: ignore[assignment]  # If the client is not installed, use the one from the SDK
+                ResolverClient,
+            )
         return ResolverClient(self)
 
     @property
@@ -112,9 +124,27 @@ class DataHubClient:
         return LineageClient(self)
 
     @property
-    def assertions(self) -> AssertionsClient:  # type: ignore[return-value]  # Type is not available if assertion_client is not installed
-        if AssertionsClient is None:
-            raise SdkUsageError(
-                "AssertionsClient is not installed, please install it with `pip install acryl-datahub-cloud`"
-            )
+    def assertions(self):  # type: ignore[report-untyped-call]  # Not available due to circular import issues
+        try:
+            from acryl_datahub_cloud.sdk import AssertionsClient
+        except ImportError as e:
+            if "acryl_datahub_cloud" in str(e):
+                raise SdkUsageError(
+                    "AssertionsClient is not installed, please install it with `pip install acryl-datahub-cloud`"
+                ) from e
+            else:
+                raise e
         return AssertionsClient(self)
+
+    @property
+    def subscriptions(self):  # type: ignore[report-untyped-call]  # Not available due to circular import issues
+        try:
+            from acryl_datahub_cloud.sdk import SubscriptionClient
+        except ImportError as e:
+            if "acryl_datahub_cloud" in str(e):
+                raise SdkUsageError(
+                    "SubscriptionClient is not installed, please install it with `pip install acryl-datahub-cloud`"
+                ) from e
+            else:
+                raise e
+        return SubscriptionClient(self)

datahub/sdk/mlmodel.py

@@ -24,6 +24,7 @@ from datahub.sdk._shared import (
     HasInstitutionalMemory,
     HasOwnership,
     HasPlatformInstance,
+    HasStructuredProperties,
     HasTags,
     HasTerms,
     HasVersion,
@@ -31,6 +32,7 @@ from datahub.sdk._shared import (
     LinksInputType,
     MLTrainingJobInputType,
     OwnersInputType,
+    StructuredPropertyInputType,
     TagsInputType,
     TermsInputType,
     TrainingMetricsInputType,
@@ -50,6 +52,7 @@ class MLModel(
     HasTerms,
     HasDomain,
     HasVersion,
+    HasStructuredProperties,
     Entity,
 ):
     __slots__ = ()
@@ -82,53 +85,43 @@ class MLModel(
         model_group: Optional[Union[str, MlModelGroupUrn]] = None,
         training_jobs: Optional[MLTrainingJobInputType] = None,
         downstream_jobs: Optional[MLTrainingJobInputType] = None,
+        structured_properties: Optional[StructuredPropertyInputType] = 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)
+        # Initialize properties in logical groups
+        self._init_basic_properties(
+            version=version,
+            name=name,
+            aliases=aliases,
+            description=description,
+            external_url=external_url,
+            custom_properties=custom_properties,
+            created=created,
+            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)
+        self._init_ml_specific_properties(
+            training_metrics=training_metrics,
+            hyper_params=hyper_params,
+            model_group=model_group,
+            training_jobs=training_jobs,
+            downstream_jobs=downstream_jobs,
+        )
+
+        self._init_metadata_properties(
+            owners=owners,
+            links=links,
+            tags=tags,
+            terms=terms,
+            domain=domain,
+            structured_properties=structured_properties,
+        )
 
     @classmethod
     def _new_from_graph(cls, urn: Urn, current_aspects: AspectBag) -> Self:
@@ -299,3 +292,73 @@ class MLModel(
             props.downstreamJobs = [
                 job for job in props.downstreamJobs if job != job_str
             ]
+
+    def _init_basic_properties(
+        self,
+        version: Optional[str] = None,
+        name: Optional[str] = None,
+        aliases: Optional[List[str]] = None,
+        description: Optional[str] = None,
+        external_url: Optional[str] = None,
+        custom_properties: Optional[Dict[str, str]] = None,
+        created: Optional[datetime] = None,
+        last_modified: Optional[datetime] = None,
+    ) -> None:
+        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 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)
+
+    def _init_ml_specific_properties(
+        self,
+        training_metrics: Optional[TrainingMetricsInputType] = None,
+        hyper_params: Optional[HyperParamsInputType] = None,
+        model_group: Optional[Union[str, MlModelGroupUrn]] = None,
+        training_jobs: Optional[MLTrainingJobInputType] = None,
+        downstream_jobs: Optional[MLTrainingJobInputType] = None,
+    ) -> None:
+        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 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)
+
+    def _init_metadata_properties(
+        self,
+        owners: Optional[OwnersInputType] = None,
+        links: Optional[LinksInputType] = None,
+        tags: Optional[TagsInputType] = None,
+        terms: Optional[TermsInputType] = None,
+        domain: Optional[DomainInputType] = None,
+        structured_properties: Optional[StructuredPropertyInputType] = None,
+    ) -> None:
+        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 structured_properties is not None:
+            for key, value in structured_properties.items():
+                self.set_structured_property(property_urn=key, values=value)

datahub/sdk/mlmodelgroup.py

@@ -17,10 +17,12 @@ from datahub.sdk._shared import (
     HasInstitutionalMemory,
     HasOwnership,
     HasPlatformInstance,
+    HasStructuredProperties,
     HasTags,
     HasTerms,
     LinksInputType,
     OwnersInputType,
+    StructuredPropertyInputType,
     TagsInputType,
     TermsInputType,
     make_time_stamp,
@@ -36,6 +38,7 @@ class MLModelGroup(
     HasTags,
     HasTerms,
     HasDomain,
+    HasStructuredProperties,
     Entity,
 ):
     __slots__ = ()
@@ -66,6 +69,7 @@ class MLModelGroup(
         domain: Optional[DomainInputType] = None,
         training_jobs: Optional[Sequence[Union[str, DataProcessInstanceUrn]]] = None,
         downstream_jobs: Optional[Sequence[Union[str, DataProcessInstanceUrn]]] = None,
+        structured_properties: Optional[StructuredPropertyInputType] = None,
         extra_aspects: ExtraAspectsType = None,
     ):
         urn = MlModelGroupUrn(platform=platform, name=id, env=env)
@@ -105,6 +109,9 @@ class MLModelGroup(
             self.set_training_jobs(training_jobs)
         if downstream_jobs is not None:
             self.set_downstream_jobs(downstream_jobs)
+        if structured_properties is not None:
+            for key, value in structured_properties.items():
+                self.set_structured_property(property_urn=key, values=value)
 
     @classmethod
     def _new_from_graph(cls, urn: Urn, current_aspects: AspectBag) -> Self:

datahub/sdk/search_client.py

@@ -19,6 +19,7 @@ from datahub.sdk.search_filters import (
     _OrFilters,
     _StatusFilter,
 )
+from datahub.utilities.ordered_set import OrderedSet
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
@@ -80,7 +81,7 @@ def compute_entity_types(
 ) -> Optional[List[str]]:
     found_filters = False
     found_positive_filters = False
-    entity_types: List[str] = []
+    entity_types: OrderedSet[str] = OrderedSet()
     for ands in filters:
         for clause in ands["and"]:
             if clause.field == _EntityTypeFilter.ENTITY_TYPE_FIELD:
@@ -88,7 +89,7 @@ def compute_entity_types(
                 if not clause.negated:
                     found_positive_filters = True
 
-                entity_types.extend(clause.values)
+                entity_types.update(clause.values)
 
     if not found_filters:
         # If we didn't find any filters, use None so we use the default set.
@@ -100,7 +101,7 @@ def compute_entity_types(
         # still want to use the default set.
         return None
 
-    return entity_types
+    return list(entity_types)
 
 
 class SearchClient:

datahub/sdk/search_filters.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import abc
 from typing import (
+    TYPE_CHECKING,
+    Annotated,
     Any,
     ClassVar,
     Iterator,
@@ -15,7 +17,10 @@ from typing import (
 import pydantic
 
 from datahub.configuration.common import ConfigModel
-from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
+from datahub.configuration.pydantic_migration_helpers import (
+    PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR,
+    PYDANTIC_VERSION_2,
+)
 from datahub.ingestion.graph.client import flexible_entity_type_to_graphql
 from datahub.ingestion.graph.filters import (
     FilterOperator,
@@ -42,12 +47,29 @@ class _BaseFilter(ConfigModel):
             populate_by_name = True
 
     @abc.abstractmethod
-    def compile(self) -> _OrFilters:
-        pass
+    def compile(self) -> _OrFilters: ...
 
     def dfs(self) -> Iterator[_BaseFilter]:
         yield self
 
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        if cls is _BaseFilter:
+            raise ValueError("Cannot get discriminator for _BaseFilter")
+        if PYDANTIC_VERSION_2:
+            fields: dict = cls.model_fields  # type: ignore
+        else:
+            fields = cls.__fields__  # type: ignore
+
+        # Assumes that there's only one field name per filter.
+        # If that's not the case, this method should be overridden.
+        if len(fields.keys()) != 1:
+            raise ValueError(
+                f"Found multiple fields that could be the discriminator for this filter: {list(fields.keys())}"
+            )
+        name, field = next(iter(fields.items()))
+        return field.alias or name  # type: ignore
+
 
 class _EntityTypeFilter(_BaseFilter):
     """Filter for specific entity types.
@@ -74,15 +96,19 @@ class _EntityTypeFilter(_BaseFilter):
 
 
 class _EntitySubtypeFilter(_BaseFilter):
-    entity_subtype: str = pydantic.Field(
+    entity_subtype: List[str] = pydantic.Field(
         description="The entity subtype to filter on. Can be 'Table', 'View', 'Source', etc. depending on the native platform's concepts.",
     )
 
+    @pydantic.validator("entity_subtype", pre=True)
+    def validate_entity_subtype(cls, v: str) -> List[str]:
+        return [v] if not isinstance(v, list) else v
+
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
             field="typeNames",
             condition="EQUAL",
-            values=[self.entity_subtype],
+            values=self.entity_subtype,
         )
 
     def compile(self) -> _OrFilters:
@@ -196,6 +222,10 @@ class _CustomCondition(_BaseFilter):
         )
         return [{"and": [rule]}]
 
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        return "_custom"
+
 
 class _And(_BaseFilter):
     """Represents an AND conjunction of filters."""
@@ -302,31 +332,69 @@ class _Not(_BaseFilter):
         yield from self.not_.dfs()
 
 
-# TODO: With pydantic 2, we can use a RootModel with a
-# discriminated union to make the error messages more informative.
-Filter = Union[
-    _And,
-    _Or,
-    _Not,
-    _EntityTypeFilter,
-    _EntitySubtypeFilter,
-    _StatusFilter,
-    _PlatformFilter,
-    _DomainFilter,
-    _EnvFilter,
-    _CustomCondition,
-]
-
-
-# Required to resolve forward references to "Filter"
-if PYDANTIC_VERSION_2:
-    _And.model_rebuild()  # type: ignore
-    _Or.model_rebuild()  # type: ignore
-    _Not.model_rebuild()  # type: ignore
-else:
+def _filter_discriminator(v: Any) -> Optional[str]:
+    if isinstance(v, _BaseFilter):
+        return v._field_discriminator()
+
+    if not isinstance(v, dict):
+        return None
+
+    keys = list(v.keys())
+    if len(keys) == 1:
+        return keys[0]
+    elif set(keys).issuperset({"field", "condition"}):
+        return _CustomCondition._field_discriminator()
+
+    return None
+
+
+if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
+    # The `not TYPE_CHECKING` bit is required to make the linter happy,
+    # since we currently only run mypy with pydantic v1.
+    Filter = Union[
+        _And,
+        _Or,
+        _Not,
+        _EntityTypeFilter,
+        _EntitySubtypeFilter,
+        _StatusFilter,
+        _PlatformFilter,
+        _DomainFilter,
+        _EnvFilter,
+        _CustomCondition,
+    ]
+
     _And.update_forward_refs()
     _Or.update_forward_refs()
     _Not.update_forward_refs()
+else:
+    from pydantic import Discriminator, Tag
+
+    # TODO: Once we're fully on pydantic 2, we can use a RootModel here.
+    # That way we'd be able to attach methods to the Filter type.
+    # e.g. replace load_filters(...) with Filter.load(...)
+    Filter = Annotated[
+        Union[
+            Annotated[_And, Tag(_And._field_discriminator())],
+            Annotated[_Or, Tag(_Or._field_discriminator())],
+            Annotated[_Not, Tag(_Not._field_discriminator())],
+            Annotated[_EntityTypeFilter, Tag(_EntityTypeFilter._field_discriminator())],
+            Annotated[
+                _EntitySubtypeFilter, Tag(_EntitySubtypeFilter._field_discriminator())
+            ],
+            Annotated[_StatusFilter, Tag(_StatusFilter._field_discriminator())],
+            Annotated[_PlatformFilter, Tag(_PlatformFilter._field_discriminator())],
+            Annotated[_DomainFilter, Tag(_DomainFilter._field_discriminator())],
+            Annotated[_EnvFilter, Tag(_EnvFilter._field_discriminator())],
+            Annotated[_CustomCondition, Tag(_CustomCondition._field_discriminator())],
+        ],
+        Discriminator(_filter_discriminator),
+    ]
+
+    # Required to resolve forward references to "Filter"
+    _And.model_rebuild()  # type: ignore
+    _Or.model_rebuild()  # type: ignore
+    _Not.model_rebuild()  # type: ignore
 
 
 def load_filters(obj: Any) -> Filter:

datahub/specific/chart.py

@@ -77,7 +77,7 @@ class ChartPatchBuilder(
             ChartInfo.ASPECT_NAME,
             "add",
             path=("inputEdges", input_urn),
-            value=input_urn,
+            value=input_edge,
         )
         return self
 

datahub/specific/dataproduct.py

@@ -9,6 +9,9 @@ from datahub.metadata.schema_classes import (
 )
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
+from datahub.specific.aspect_helpers.structured_properties import (
+    HasStructuredPropertiesPatch,
+)
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
 from datahub.specific.aspect_helpers.terms import HasTermsPatch
 
@@ -16,6 +19,7 @@ from datahub.specific.aspect_helpers.terms import HasTermsPatch
 class DataProductPatchBuilder(
     HasOwnershipPatch,
     HasCustomPropertiesPatch,
+    HasStructuredPropertiesPatch,
     HasTagsPatch,
     HasTermsPatch,
     MetadataPatchProposal,