acryl-datahub 1.2.0.3rc1__py3-none-any.whl → 1.2.0.4__py3-none-any.whl

datahub/metadata/schemas/DatasetUsageStatistics.avsc

@@ -130,6 +130,10 @@
       "doc": "The optional messageId, if provided serves as a custom user-defined unique identifier for an aspect value."
     },
     {
+      "Searchable": {
+        "fieldType": "COUNT",
+        "hasValuesFieldName": "hasUniqueUserCount"
+      },
       "TimeseriesField": {},
       "type": [
         "null",
@@ -140,6 +144,10 @@
       "doc": "Unique user count"
     },
     {
+      "Searchable": {
+        "fieldType": "COUNT",
+        "hasValuesFieldName": "hasTotalSqlQueriesCount"
+      },
       "TimeseriesField": {},
       "type": [
         "null",

datahub/sdk/entity_client.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Union, overload
+from typing import TYPE_CHECKING, Optional, Union, overload
 
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.emitter.rest_emitter import EmitMode
 from datahub.errors import IngestionAttributionWarning, ItemNotFoundError, SdkUsageError
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import (
@@ -133,7 +134,7 @@ class EntityClient:
 
         return entity
 
-    def create(self, entity: Entity) -> None:
+    def create(self, entity: Entity, *, emit_mode: Optional[EmitMode] = None) -> None:
         mcps = []
 
         if self._graph.exists(str(entity.urn)):
@@ -152,9 +153,12 @@ class EntityClient:
         )
         mcps.extend(entity.as_mcps(models.ChangeTypeClass.CREATE))
 
-        self._graph.emit_mcps(mcps)
+        if emit_mode:
+            self._graph.emit_mcps(mcps, emit_mode=emit_mode)
+        else:
+            self._graph.emit_mcps(mcps)
 
-    def upsert(self, entity: Entity) -> None:
+    def upsert(self, entity: Entity, *, emit_mode: Optional[EmitMode] = None) -> None:
         if entity._prev_aspects is None and self._graph.exists(str(entity.urn)):
             warnings.warn(
                 f"The entity {entity.urn} already exists. This operation will partially overwrite the existing entity.",
@@ -164,9 +168,17 @@ class EntityClient:
             # TODO: If there are no previous aspects but the entity exists, should we delete aspects that are not present here?
 
         mcps = entity.as_mcps(models.ChangeTypeClass.UPSERT)
-        self._graph.emit_mcps(mcps)
+        if emit_mode:
+            self._graph.emit_mcps(mcps, emit_mode=emit_mode)
+        else:
+            self._graph.emit_mcps(mcps)
 
-    def update(self, entity: Union[Entity, MetadataPatchProposal]) -> None:
+    def update(
+        self,
+        entity: Union[Entity, MetadataPatchProposal],
+        *,
+        emit_mode: Optional[EmitMode] = None,
+    ) -> None:
         if isinstance(entity, MetadataPatchProposal):
             return self._update_patch(entity)
 
@@ -179,7 +191,10 @@ class EntityClient:
         # -> probably add a "mode" parameter that can be "update" (e.g. if not modified) or "update_force"
 
         mcps = entity.as_mcps(models.ChangeTypeClass.UPSERT)
-        self._graph.emit_mcps(mcps)
+        if emit_mode:
+            self._graph.emit_mcps(mcps, emit_mode=emit_mode)
+        else:
+            self._graph.emit_mcps(mcps)
 
     def _update_patch(
         self, updater: MetadataPatchProposal, check_exists: bool = True

datahub/sdk/search_client.py

@@ -112,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)
@@ -120,5 +122,6 @@ class SearchClient:
             status=None,
             extra_or_filters=compiled_filters,
             entity_types=types,
+            skip_cache=skip_cache,
         ):
             yield Urn.from_string(urn)

datahub/specific/aspect_helpers/fine_grained_lineage.py

@@ -0,0 +1,76 @@
+from abc import abstractmethod
+from typing import List, Tuple
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
+from datahub.metadata.schema_classes import (
+    FineGrainedLineageClass as FineGrainedLineage,
+)
+
+
+class HasFineGrainedLineagePatch(MetadataPatchProposal):
+    @abstractmethod
+    def _fine_grained_lineage_location(self) -> Tuple[str, PatchPath]:
+        """Return the aspect name where fine-grained lineage is stored."""
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    @staticmethod
+    def _get_fine_grained_key(
+        fine_grained_lineage: FineGrainedLineage,
+    ) -> Tuple[str, str, str]:
+        downstreams = fine_grained_lineage.downstreams or []
+        if len(downstreams) != 1:
+            raise TypeError("Cannot patch with more or less than one downstream.")
+        transform_op = fine_grained_lineage.transformOperation or "NONE"
+        downstream_urn = downstreams[0]
+        query_id = fine_grained_lineage.query or "NONE"
+        return transform_op, downstream_urn, query_id
+
+    def add_fine_grained_lineage(
+        self, fine_grained_lineage: FineGrainedLineage
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        (
+            transform_op,
+            downstream_urn,
+            query_id,
+        ) = self._get_fine_grained_key(fine_grained_lineage)
+        for upstream_urn in fine_grained_lineage.upstreams or []:
+            self._add_patch(
+                aspect_name,
+                "add",
+                path=(*path, transform_op, downstream_urn, query_id, upstream_urn),
+                value={"confidenceScore": fine_grained_lineage.confidenceScore},
+            )
+        return self
+
+    def remove_fine_grained_lineage(
+        self, fine_grained_lineage: FineGrainedLineage
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        (
+            transform_op,
+            downstream_urn,
+            query_id,
+        ) = self._get_fine_grained_key(fine_grained_lineage)
+        for upstream_urn in fine_grained_lineage.upstreams or []:
+            self._add_patch(
+                aspect_name,
+                "remove",
+                path=(*path, transform_op, downstream_urn, query_id, upstream_urn),
+                value={},
+            )
+        return self
+
+    def set_fine_grained_lineages(
+        self, fine_grained_lineages: List[FineGrainedLineage]
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        self._add_patch(
+            aspect_name,
+            "add",
+            path=path,
+            value=fine_grained_lineages,
+        )
+        return self

datahub/specific/datajob.py

@@ -1,15 +1,19 @@
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Set, Tuple, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
 from datahub.metadata.schema_classes import (
     DataJobInfoClass as DataJobInfo,
     DataJobInputOutputClass as DataJobInputOutput,
     EdgeClass as Edge,
+    FineGrainedLineageClass as FineGrainedLineage,
     KafkaAuditHeaderClass,
     SystemMetadataClass,
 )
 from datahub.metadata.urns import SchemaFieldUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
+from datahub.specific.aspect_helpers.fine_grained_lineage import (
+    HasFineGrainedLineagePatch,
+)
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
 from datahub.specific.aspect_helpers.terms import HasTermsPatch
@@ -20,6 +24,7 @@ class DataJobPatchBuilder(
     HasCustomPropertiesPatch,
     HasTagsPatch,
     HasTermsPatch,
+    HasFineGrainedLineagePatch,
     MetadataPatchProposal,
 ):
     def __init__(
@@ -40,10 +45,19 @@ class DataJobPatchBuilder(
             urn, system_metadata=system_metadata, audit_header=audit_header
         )
 
+        # Track fine-grained lineages for DataJob-specific handling
+        self._fine_grained_lineages_to_add: List[FineGrainedLineage] = []
+        self._fine_grained_lineage_keys_to_remove: Set[Tuple[str, str, str]] = set()
+        self._fine_grained_lineages_set: Optional[List[FineGrainedLineage]] = None
+
     @classmethod
     def _custom_properties_location(cls) -> Tuple[str, PatchPath]:
         return DataJobInfo.ASPECT_NAME, ("customProperties",)
 
+    @classmethod
+    def _fine_grained_lineage_location(cls) -> Tuple[str, PatchPath]:
+        return DataJobInputOutput.ASPECT_NAME, ("fineGrainedLineages",)
+
     def add_input_datajob(self, input: Union[Edge, Urn, str]) -> "DataJobPatchBuilder":
         """
         Adds an input data job to the DataJobPatchBuilder.

datahub/specific/dataset.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Generic, List, Optional, Tuple, TypeVar, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
@@ -17,6 +18,9 @@ from datahub.metadata.schema_classes import (
 )
 from datahub.metadata.urns import DatasetUrn, TagUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
+from datahub.specific.aspect_helpers.fine_grained_lineage import (
+    HasFineGrainedLineagePatch,
+)
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.structured_properties import (
     HasStructuredPropertiesPatch,
@@ -99,6 +103,7 @@ class DatasetPatchBuilder(
     HasStructuredPropertiesPatch,
     HasTagsPatch,
     HasTermsPatch,
+    HasFineGrainedLineagePatch,
     MetadataPatchProposal,
 ):
     def __init__(
@@ -115,6 +120,10 @@ class DatasetPatchBuilder(
     def _custom_properties_location(cls) -> Tuple[str, PatchPath]:
         return DatasetProperties.ASPECT_NAME, ("customProperties",)
 
+    @classmethod
+    def _fine_grained_lineage_location(cls) -> Tuple[str, PatchPath]:
+        return UpstreamLineage.ASPECT_NAME, ("fineGrainedLineages",)
+
     def add_upstream_lineage(self, upstream: Upstream) -> "DatasetPatchBuilder":
         self._add_patch(
             UpstreamLineage.ASPECT_NAME,
@@ -144,75 +153,44 @@ class DatasetPatchBuilder(
     def add_fine_grained_upstream_lineage(
         self, fine_grained_lineage: FineGrainedLineage
     ) -> "DatasetPatchBuilder":
-        (
-            transform_op,
-            downstream_urn,
-            query_id,
-        ) = DatasetPatchBuilder.get_fine_grained_key(fine_grained_lineage)
-        for upstream_urn in fine_grained_lineage.upstreams or []:
-            self._add_patch(
-                UpstreamLineage.ASPECT_NAME,
-                "add",
-                path=self._build_fine_grained_path(
-                    transform_op, downstream_urn, query_id, upstream_urn
-                ),
-                value={"confidenceScore": fine_grained_lineage.confidenceScore},
-            )
-        return self
-
-    @staticmethod
-    def get_fine_grained_key(
-        fine_grained_lineage: FineGrainedLineage,
-    ) -> Tuple[str, str, str]:
-        downstreams = fine_grained_lineage.downstreams or []
-        if len(downstreams) != 1:
-            raise TypeError("Cannot patch with more or less than one downstream.")
-        transform_op = fine_grained_lineage.transformOperation or "NONE"
-        downstream_urn = downstreams[0]
-        query_id = fine_grained_lineage.query or "NONE"
-        return transform_op, downstream_urn, query_id
-
-    @classmethod
-    def _build_fine_grained_path(
-        cls, transform_op: str, downstream_urn: str, query_id: str, upstream_urn: str
-    ) -> PatchPath:
-        return (
-            "fineGrainedLineages",
-            transform_op,
-            downstream_urn,
-            query_id,
-            upstream_urn,
+        """
+        Deprecated: Use `add_fine_grained_lineage` instead.
+        """
+        warnings.warn(
+            "add_fine_grained_upstream_lineage() is deprecated."
+            " Use add_fine_grained_lineage() instead.",
+            DeprecationWarning,
+            stacklevel=2,
         )
+        return self.add_fine_grained_lineage(fine_grained_lineage)
 
     def remove_fine_grained_upstream_lineage(
         self, fine_grained_lineage: FineGrainedLineage
     ) -> "DatasetPatchBuilder":
-        (
-            transform_op,
-            downstream_urn,
-            query_id,
-        ) = DatasetPatchBuilder.get_fine_grained_key(fine_grained_lineage)
-        for upstream_urn in fine_grained_lineage.upstreams or []:
-            self._add_patch(
-                UpstreamLineage.ASPECT_NAME,
-                "remove",
-                path=self._build_fine_grained_path(
-                    transform_op, downstream_urn, query_id, upstream_urn
-                ),
-                value={},
-            )
-        return self
+        """
+        Deprecated: Use `remove_fine_grained_lineage` instead.
+        """
+        warnings.warn(
+            "remove_fine_grained_upstream_lineage() is deprecated."
+            " Use remove_fine_grained_lineage() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.remove_fine_grained_lineage(fine_grained_lineage)
 
     def set_fine_grained_upstream_lineages(
         self, fine_grained_lineages: List[FineGrainedLineage]
     ) -> "DatasetPatchBuilder":
-        self._add_patch(
-            UpstreamLineage.ASPECT_NAME,
-            "add",
-            path=("fineGrainedLineages",),
-            value=fine_grained_lineages,
+        """
+        Deprecated: Use `set_fine_grained_lineages` instead.
+        """
+        warnings.warn(
+            "set_fine_grained_upstream_lineages() is deprecated."
+            " Use set_fine_grained_lineages() instead.",
+            DeprecationWarning,
+            stacklevel=2,
         )
-        return self
+        return self.set_fine_grained_lineages(fine_grained_lineages)
 
     def for_field(
         self, field_path: str, editable: bool = True

datahub/utilities/mapping.py

@@ -83,7 +83,7 @@ class Constants:
     MATCH = "match"
     USER_OWNER = "user"
     GROUP_OWNER = "group"
-    OPERAND_DATATYPE_SUPPORTED = [int, bool, str, float]
+    OPERAND_DATATYPE_SUPPORTED = [int, bool, str, float, list]
     TAG_PARTITION_KEY = "PARTITION_KEY"
     TAG_DIST_KEY = "DIST_KEY"
     TAG_SORT_KEY = "SORT_KEY"
@@ -455,7 +455,34 @@ class OperationProcessor:
         # function to check if a match clause is satisfied to a value.
         if not any(
             isinstance(raw_props_value, t) for t in Constants.OPERAND_DATATYPE_SUPPORTED
-        ) or not isinstance(raw_props_value, type(match_clause)):
+        ):
+            return None
+
+        # Handle list values by checking if any item in the list matches
+        if isinstance(raw_props_value, list):
+            # For lists, we need to find at least one matching item
+            # Return a match with the concatenated values of all matching items
+            matching_items = []
+            for item in raw_props_value:
+                if isinstance(item, str):
+                    match = re.match(match_clause, item)
+                    if match:
+                        matching_items.append(item)
+                elif isinstance(match_clause, type(item)):
+                    match = re.match(str(match_clause), str(item))
+                    if match:
+                        matching_items.append(str(item))
+
+            if matching_items:
+                # Create a synthetic match object with all matching items joined
+                combined_value = ",".join(matching_items)
+                return re.match(
+                    ".*", combined_value
+                )  # Always matches, returns combined value
+            return None
+
+        # Handle scalar values (existing logic)
+        elif not isinstance(raw_props_value, type(match_clause)):
             return None
         elif isinstance(raw_props_value, str):
             return re.match(match_clause, raw_props_value)

datahub/utilities/server_config_util.py

@@ -242,7 +242,8 @@ class RestServiceConfig:
 
         # Check if this is a config-based feature
         if feature in config_based_features:
-            return config_based_features[feature]()
+            result = config_based_features[feature]()
+            return bool(result) if result is not None else False
 
         # For environment-based features, determine requirements based on cloud vs. non-cloud
         deployment_type = "cloud" if self.is_datahub_cloud else "core"