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

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

datahub/sql_parsing/split_statements.py

@@ -52,6 +52,7 @@ class ParserState(Enum):
     STRING = 2
     COMMENT = 3
     MULTILINE_COMMENT = 4
+    BRACKETED_IDENTIFIER = 5
 
 
 class _StatementSplitter:
@@ -141,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)
@@ -172,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":

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -4,7 +4,6 @@ import enum
 import functools
 import json
 import logging
-import os
 import pathlib
 import tempfile
 import uuid
@@ -14,10 +13,10 @@ from typing import Callable, Dict, Iterable, List, Optional, Set, Union, cast
 
 import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
+from datahub.configuration.env_vars import get_sql_agg_query_log
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.emitter.mce_builder import get_sys_time, make_ts_millis
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.sql_parsing_builder import compute_upstream_fields
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.report import Report
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -49,6 +48,7 @@ from datahub.sql_parsing.sqlglot_lineage import (
     sqlglot_lineage,
 )
 from datahub.sql_parsing.sqlglot_utils import (
+    DialectOrStr,
     _parse_statement,
     get_query_fingerprint,
     try_format_query,
@@ -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,
@@ -82,7 +83,7 @@ class QueryLogSetting(enum.Enum):
 _DEFAULT_USER_URN = CorpUserUrn("_ingestion")
 _MISSING_SESSION_ID = "__MISSING_SESSION_ID"
 _DEFAULT_QUERY_LOG_SETTING = QueryLogSetting[
-    os.getenv("DATAHUB_SQL_AGG_QUERY_LOG") or QueryLogSetting.DISABLED.name
+    get_sql_agg_query_log() or QueryLogSetting.DISABLED.name
 ]
 MAX_UPSTREAM_TABLES_COUNT = 300
 MAX_FINEGRAINEDLINEAGE_COUNT = 2000
@@ -108,6 +109,7 @@ class ObservedQuery:
     default_schema: Optional[str] = None
     query_hash: Optional[str] = None
     usage_multiplier: int = 1
+    override_dialect: Optional[DialectOrStr] = None
 
     # Use this to store additional key-value information about the query for debugging.
     extra_info: Optional[dict] = None
@@ -140,6 +142,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:
@@ -188,6 +191,7 @@ class QueryMetadata:
             source=models.QuerySourceClass.SYSTEM,
             created=self.make_created_audit_stamp(),
             lastModified=self.make_last_modified_audit_stamp(),
+            origin=self.origin.urn() if self.origin else None,
         )
 
 
@@ -263,7 +267,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
 
@@ -629,6 +633,9 @@ class SqlParsingAggregator(Closeable):
             TableSwap,
         ],
     ) -> None:
+        """
+        This assumes that queries come in order of increasing timestamps.
+        """
         if isinstance(item, KnownQueryLineageInfo):
             self.add_known_query_lineage(item)
         elif isinstance(item, KnownLineageMapping):
@@ -831,6 +838,7 @@ class SqlParsingAggregator(Closeable):
             session_id=session_id,
             timestamp=observed.timestamp,
             user=observed.user,
+            override_dialect=observed.override_dialect,
         )
         if parsed.debug_info.error:
             self.report.observed_query_parse_failures.append(
@@ -859,7 +867,7 @@ class SqlParsingAggregator(Closeable):
                 downstream=parsed.out_tables[0] if parsed.out_tables else None,
                 column_lineage=parsed.column_lineage,
                 # TODO: We need a full list of columns referenced, not just the out tables.
-                column_usage=compute_upstream_fields(parsed),
+                column_usage=self._compute_upstream_fields(parsed),
                 inferred_schema=infer_output_schema(parsed),
                 confidence_score=parsed.debug_info.confidence,
                 extra_info=observed.extra_info,
@@ -948,6 +956,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,
             )
         )
@@ -1147,7 +1156,7 @@ class SqlParsingAggregator(Closeable):
                 actor=None,
                 upstreams=parsed.in_tables,
                 column_lineage=parsed.column_lineage or [],
-                column_usage=compute_upstream_fields(parsed),
+                column_usage=self._compute_upstream_fields(parsed),
                 confidence_score=parsed.debug_info.confidence,
             )
         )
@@ -1164,6 +1173,7 @@ class SqlParsingAggregator(Closeable):
         session_id: str = _MISSING_SESSION_ID,
         timestamp: Optional[datetime] = None,
         user: Optional[Union[CorpUserUrn, CorpGroupUrn]] = None,
+        override_dialect: Optional[DialectOrStr] = None,
     ) -> SqlParsingResult:
         with self.report.sql_parsing_timer:
             parsed = sqlglot_lineage(
@@ -1171,6 +1181,7 @@ class SqlParsingAggregator(Closeable):
                 schema_resolver=schema_resolver,
                 default_db=default_db,
                 default_schema=default_schema,
+                override_dialect=override_dialect,
             )
         self.report.num_sql_parsed += 1
 
@@ -1329,11 +1340,25 @@ class SqlParsingAggregator(Closeable):
                 upstreams.setdefault(upstream, query.query_id)
 
             for lineage_info in query.column_lineage:
-                for upstream_ref in lineage_info.upstreams:
-                    cll[lineage_info.downstream.column].setdefault(
-                        SchemaFieldUrn(upstream_ref.table, upstream_ref.column),
-                        query.query_id,
+                if (
+                    not lineage_info.downstream.column
+                    or not lineage_info.downstream.column.strip()
+                ):
+                    logger.debug(
+                        f"Skipping lineage entry with empty downstream column in query {query.query_id}"
                     )
+                    continue
+
+                for upstream_ref in lineage_info.upstreams:
+                    if upstream_ref.column and upstream_ref.column.strip():
+                        cll[lineage_info.downstream.column].setdefault(
+                            SchemaFieldUrn(upstream_ref.table, upstream_ref.column),
+                            query.query_id,
+                        )
+                    else:
+                        logger.debug(
+                            f"Skipping empty column reference in lineage for query {query.query_id}"
+                        )
 
         # Finally, we can build our lineage edge.
         required_queries = OrderedSet[QueryId]()
@@ -1491,9 +1516,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 +1599,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 +1640,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 +1700,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,20 +1740,30 @@ 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,
         )
 
         return resolved_query
 
+    @staticmethod
+    def _compute_upstream_fields(
+        result: SqlParsingResult,
+    ) -> Dict[UrnStr, Set[UrnStr]]:
+        upstream_fields: Dict[UrnStr, Set[UrnStr]] = defaultdict(set)
+        for cl in result.column_lineage or []:
+            for upstream in cl.upstreams:
+                upstream_fields[upstream.table].add(upstream.column)
+        return upstream_fields
+
     def _gen_usage_statistics_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
         if not self._usage_aggregator:
             return