acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/sdk/main_client.py

@@ -1,21 +1,16 @@
 from __future__ import annotations
 
-from typing import Optional, overload
+from typing import TYPE_CHECKING, Optional, overload
 
 from datahub.errors import SdkUsageError
 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
+if TYPE_CHECKING:
+    from datahub.sdk.resolver_client import ResolverClient
 
 
 class DataHubClient:
@@ -74,7 +69,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 +84,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 +95,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 +107,16 @@ class DataHubClient:
         return EntityClient(self)
 
     @property
-    def resolve(self) -> ResolverClient:
+    def resolve(self) -> "ResolverClient":
+        try:
+            from acryl_datahub_cloud.sdk import (  # type: ignore[import-not-found]
+                ResolverClient,
+            )
+        except ImportError:
+            # If the client is not installed, use the one from the SDK.
+            from datahub.sdk.resolver_client import (  # type: ignore[assignment]
+                ResolverClient,
+            )
         return ResolverClient(self)
 
     @property
@@ -112,9 +128,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,92 @@ class MLModel(
             props.downstreamJobs = [
                 job for job in props.downstreamJobs if job != job_str
             ]
+
+    @property
+    def deployments(self) -> Optional[List[str]]:
+        return self._ensure_model_props().deployments
+
+    def set_deployments(self, deployments: Sequence[str]) -> None:
+        self._ensure_model_props().deployments = list(deployments)
+
+    def add_deployment(self, deployment: str) -> None:
+        props = self._ensure_model_props()
+        if props.deployments is None:
+            props.deployments = []
+        if deployment not in props.deployments:
+            props.deployments.append(deployment)
+
+    def remove_deployment(self, deployment: str) -> None:
+        props = self._ensure_model_props()
+        if props.deployments is not None:
+            props.deployments = [d for d in props.deployments if d != deployment]
+
+    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:
@@ -111,6 +112,8 @@ class SearchClient:
         self,
         query: Optional[str] = None,
         filter: Optional[Filter] = None,
+        *,
+        skip_cache: bool = False,
     ) -> Iterable[Urn]:
         # TODO: Add better limit / pagination support.
         types, compiled_filters = compile_filters(filter)
@@ -119,5 +122,6 @@ class SearchClient:
             status=None,
             extra_or_filters=compiled_filters,
             entity_types=types,
+            skip_cache=skip_cache,
         ):
             yield Urn.from_string(urn)