acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

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/aspect_helpers/siblings.py

@@ -0,0 +1,73 @@
+from typing import List
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
+
+
+class HasSiblingsPatch(MetadataPatchProposal):
+    def add_sibling(self, sibling_urn: str, primary: bool = False) -> Self:
+        """Add a sibling relationship to the entity.
+
+        Args:
+            sibling_urn: The URN of the sibling entity to add.
+            primary: Whether this entity should be marked as primary in the relationship.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME,
+            "add",
+            path=("siblings", sibling_urn),
+            value=sibling_urn,
+        )
+
+        # Set primary flag if specified
+        if primary:
+            self._add_patch(
+                Siblings.ASPECT_NAME,
+                "add",
+                path=("primary",),
+                value=primary,
+            )
+
+        return self
+
+    def remove_sibling(self, sibling_urn: str) -> Self:
+        """Remove a sibling relationship from the entity.
+
+        Args:
+            sibling_urn: The URN of the sibling entity to remove.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME,
+            "remove",
+            path=("siblings", sibling_urn),
+            value={},
+        )
+        return self
+
+    def set_siblings(self, sibling_urns: List[str], primary: bool = False) -> Self:
+        """Set the complete list of siblings for the entity.
+
+        This will replace all existing siblings with the new list.
+
+        Args:
+            sibling_urns: The list of sibling URNs to set.
+            primary: Whether this entity should be marked as primary.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME, "add", path=("siblings",), value=sibling_urns
+        )
+
+        self._add_patch(Siblings.ASPECT_NAME, "add", path=("primary",), value=primary)
+
+        return self

datahub/specific/aspect_helpers/structured_properties.py

@@ -70,3 +70,30 @@ class HasStructuredPropertiesPatch(MetadataPatchProposal):
             ),
         )
         return self
+
+    def set_structured_property_manual(
+        self, property: StructuredPropertyValueAssignmentClass
+    ) -> Self:
+        """Add or update a structured property, using a StructuredPropertyValueAssignmentClass object."""
+
+        self.remove_structured_property(property.propertyUrn)
+        self._add_patch(
+            StructuredPropertiesClass.ASPECT_NAME,
+            "add",
+            path=("properties", property.propertyUrn),
+            value=property,
+        )
+        return self
+
+    def add_structured_property_manual(
+        self, property: StructuredPropertyValueAssignmentClass
+    ) -> Self:
+        """Add a structured property, using a StructuredPropertyValueAssignmentClass object."""
+
+        self._add_patch(
+            StructuredPropertiesClass.ASPECT_NAME,
+            "add",
+            path=("properties", property.propertyUrn),
+            value=property,
+        )
+        return self

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/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/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/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,7 +18,11 @@ 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.siblings import HasSiblingsPatch
 from datahub.specific.aspect_helpers.structured_properties import (
     HasStructuredPropertiesPatch,
 )
@@ -99,6 +104,8 @@ class DatasetPatchBuilder(
     HasStructuredPropertiesPatch,
     HasTagsPatch,
     HasTermsPatch,
+    HasFineGrainedLineagePatch,
+    HasSiblingsPatch,
     MetadataPatchProposal,
 ):
     def __init__(
@@ -115,6 +122,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 +155,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
@@ -292,3 +272,15 @@ class DatasetPatchBuilder(
                 value=timestamp,
             )
         return self
+
+    def set_external_url(
+        self, external_url: Optional[str] = None
+    ) -> "DatasetPatchBuilder":
+        if external_url is not None:
+            self._add_patch(
+                DatasetProperties.ASPECT_NAME,
+                "add",
+                path=("externalUrl",),
+                value=external_url,
+            )
+        return self

datahub/sql_parsing/_sqlglot_patch.py

@@ -163,8 +163,7 @@ def _patch_lineage() -> None:
 -    source_columns = set(find_all_in_scope(select, exp.Column))
 +    source_columns = list(find_all_in_scope(select, exp.Column))
 
--    # If the source is a UDTF find columns used in the UTDF to generate the table
-+    # If the source is a UDTF find columns used in the UDTF to generate the table
+     # If the source is a UDTF find columns used in the UDTF to generate the table
 +    source = scope.expression
      if isinstance(source, exp.UDTF):
 -        source_columns |= set(source.find_all(exp.Column))

datahub/sql_parsing/fingerprint_utils.py

@@ -0,0 +1,6 @@
+import hashlib
+
+
+def generate_hash(text: str) -> str:
+    # Once we move to Python 3.9+, we can set `usedforsecurity=False`.
+    return hashlib.sha256(text.encode("utf-8")).hexdigest()

datahub/sql_parsing/split_statements.py

@@ -1,7 +1,9 @@
+import logging
 import re
 from enum import Enum
 from typing import Iterator, List, Tuple
 
+logger = logging.getLogger(__name__)
 SELECT_KEYWORD = "SELECT"
 CASE_KEYWORD = "CASE"
 END_KEYWORD = "END"
@@ -50,6 +52,7 @@ class ParserState(Enum):
     STRING = 2
     COMMENT = 3
     MULTILINE_COMMENT = 4
+    BRACKETED_IDENTIFIER = 5
 
 
 class _StatementSplitter:
@@ -120,7 +123,9 @@ class _StatementSplitter:
         # Reset current_statement-specific state.
         self.does_select_mean_new_statement = False
         if self.current_case_statements != 0:
-            breakpoint()
+            logger.warning(
+                f"Unexpected END keyword. Current case statements: {self.current_case_statements}"
+            )
         self.current_case_statements = 0
 
     def process(self) -> Iterator[str]:
@@ -137,6 +142,10 @@ class _StatementSplitter:
                     self.state = ParserState.STRING
                     self.current_statement.append(c)
                     prev_real_char = c
+                elif c == "[":
+                    self.state = ParserState.BRACKETED_IDENTIFIER
+                    self.current_statement.append(c)
+                    prev_real_char = c
                 elif c == "-" and next_char == "-":
                     self.state = ParserState.COMMENT
                     self.current_statement.append(c)
@@ -168,6 +177,14 @@ class _StatementSplitter:
                 elif c == "'":
                     self.state = ParserState.NORMAL
 
+            elif self.state == ParserState.BRACKETED_IDENTIFIER:
+                self.current_statement.append(c)
+                if c == "]" and next_char == "]":
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                elif c == "]":
+                    self.state = ParserState.NORMAL
+
             elif self.state == ParserState.COMMENT:
                 self.current_statement.append(c)
                 if c == "\n":
@@ -233,8 +250,10 @@ class _StatementSplitter:
             ),
         )
         if (
-            is_force_new_statement_keyword and most_recent_real_char != ")"
-        ):  # usually we'd have a close paren that closes a CTE
+            is_force_new_statement_keyword
+            and not self._has_preceding_cte(most_recent_real_char)
+            and not self._is_part_of_merge_query()
+        ):
             # Force termination of current statement
             yield from self._yield_if_complete()
 
@@ -247,6 +266,14 @@ class _StatementSplitter:
         else:
             self.current_statement.append(c)
 
+    def _has_preceding_cte(self, most_recent_real_char: str) -> bool:
+        # usually we'd have a close paren that closes a CTE
+        return most_recent_real_char == ")"
+
+    def _is_part_of_merge_query(self) -> bool:
+        # In merge statement we'd have `when matched then` or `when not matched then"
+        return "".join(self.current_statement).strip().lower().endswith("then")
+
 
 def split_statements(sql: str) -> Iterator[str]:
     """