acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.2.0__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/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,

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -58,6 +58,7 @@ from datahub.sql_parsing.tool_meta_extractor import (
     ToolMetaExtractorReport,
 )
 from datahub.utilities.cooperative_timeout import CooperativeTimeoutError
+from datahub.utilities.dedup_list import deduplicate_list
 from datahub.utilities.file_backed_collections import (
     ConnectionWrapper,
     FileBackedDict,
@@ -140,6 +141,7 @@ class QueryMetadata:
 
     used_temp_tables: bool = True
 
+    extra_info: Optional[dict] = None
     origin: Optional[Urn] = None
 
     def make_created_audit_stamp(self) -> models.AuditStampClass:
@@ -263,7 +265,7 @@ class PreparsedQuery:
     query_type_props: QueryTypeProps = dataclasses.field(
         default_factory=lambda: QueryTypeProps()
     )
-    # Use this to store addtitional key-value information about query for debugging
+    # Use this to store additional key-value information about the query for debugging.
     extra_info: Optional[dict] = None
     origin: Optional[Urn] = None
 
@@ -948,6 +950,7 @@ class SqlParsingAggregator(Closeable):
                 column_usage=parsed.column_usage or {},
                 confidence_score=parsed.confidence_score,
                 used_temp_tables=session_has_temp_tables,
+                extra_info=parsed.extra_info,
                 origin=parsed.origin,
             )
         )
@@ -1491,9 +1494,9 @@ class SqlParsingAggregator(Closeable):
             return
 
         # If a query doesn't involve any allowed tables, skip it.
-        if downstream_urn is None and not any(
-            self.is_allowed_table(urn) for urn in query.upstreams
-        ):
+        if (
+            downstream_urn is None or not self.is_allowed_table(downstream_urn)
+        ) and not any(self.is_allowed_table(urn) for urn in query.upstreams):
             self.report.num_queries_skipped_due_to_filters += 1
             return
 
@@ -1574,27 +1577,33 @@ class SqlParsingAggregator(Closeable):
 
         @dataclasses.dataclass
         class QueryLineageInfo:
-            upstreams: List[UrnStr]  # this is direct upstreams, with *no temp tables*
-            column_lineage: List[ColumnLineageInfo]
+            upstreams: OrderedSet[
+                UrnStr
+            ]  # this is direct upstreams, with *no temp tables*
+            column_lineage: OrderedSet[ColumnLineageInfo]
             confidence_score: float
 
             def _merge_lineage_from(self, other_query: "QueryLineageInfo") -> None:
-                self.upstreams += other_query.upstreams
-                self.column_lineage += other_query.column_lineage
+                self.upstreams.update(other_query.upstreams)
+                self.column_lineage.update(other_query.column_lineage)
                 self.confidence_score = min(
                     self.confidence_score, other_query.confidence_score
                 )
 
+        cache: Dict[str, QueryLineageInfo] = {}
+
         def _recurse_into_query(
             query: QueryMetadata, recursion_path: List[QueryId]
         ) -> QueryLineageInfo:
             if query.query_id in recursion_path:
                 # This is a cycle, so we just return the query as-is.
                 return QueryLineageInfo(
-                    upstreams=query.upstreams,
-                    column_lineage=query.column_lineage,
+                    upstreams=OrderedSet(query.upstreams),
+                    column_lineage=OrderedSet(query.column_lineage),
                     confidence_score=query.confidence_score,
                 )
+            if query.query_id in cache:
+                return cache[query.query_id]
             recursion_path = [*recursion_path, query.query_id]
             composed_of_queries.add(query.query_id)
 
@@ -1609,7 +1618,7 @@ class SqlParsingAggregator(Closeable):
                         upstream_query = self._query_map.get(upstream_query_id)
                         if (
                             upstream_query
-                            and upstream_query.query_id not in composed_of_queries
+                            and upstream_query.query_id not in recursion_path
                         ):
                             temp_query_lineage_info = _recurse_into_query(
                                 upstream_query, recursion_path
@@ -1669,11 +1678,14 @@ class SqlParsingAggregator(Closeable):
                 ]
             )
 
-            return QueryLineageInfo(
-                upstreams=list(new_upstreams),
-                column_lineage=new_cll,
+            ret = QueryLineageInfo(
+                upstreams=new_upstreams,
+                column_lineage=OrderedSet(new_cll),
                 confidence_score=new_confidence_score,
             )
+            cache[query.query_id] = ret
+
+            return ret
 
         resolved_lineage_info = _recurse_into_query(base_query, [])
 
@@ -1706,15 +1718,15 @@ class SqlParsingAggregator(Closeable):
         )
 
         merged_query_text = ";\n\n".join(
-            [q.formatted_query_string for q in ordered_queries]
+            deduplicate_list([q.formatted_query_string for q in ordered_queries])
         )
 
         resolved_query = dataclasses.replace(
             base_query,
             query_id=composite_query_id,
             formatted_query_string=merged_query_text,
-            upstreams=resolved_lineage_info.upstreams,
-            column_lineage=resolved_lineage_info.column_lineage,
+            upstreams=list(resolved_lineage_info.upstreams),
+            column_lineage=list(resolved_lineage_info.column_lineage),
             confidence_score=resolved_lineage_info.confidence_score,
         )