acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.0.0.1rc2__py3-none-any.whl

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -30,6 +30,7 @@ from datahub.metadata.urns import (
     DatasetUrn,
     QueryUrn,
     SchemaFieldUrn,
+    Urn,
 )
 from datahub.sql_parsing.schema_resolver import (
     SchemaResolver,
@@ -139,6 +140,8 @@ class QueryMetadata:
 
     used_temp_tables: bool = True
 
+    origin: Optional[Urn] = None
+
     def make_created_audit_stamp(self) -> models.AuditStampClass:
         return models.AuditStampClass(
             time=make_ts_millis(self.latest_timestamp) or 0,
@@ -221,6 +224,7 @@ class PreparsedQuery:
     )
     # Use this to store addtitional key-value information about query for debugging
     extra_info: Optional[dict] = None
+    origin: Optional[Urn] = None
 
 
 @dataclasses.dataclass
@@ -903,6 +907,7 @@ class SqlParsingAggregator(Closeable):
                 column_usage=parsed.column_usage or {},
                 confidence_score=parsed.confidence_score,
                 used_temp_tables=session_has_temp_tables,
+                origin=parsed.origin,
             )
         )
 
@@ -1464,6 +1469,7 @@ class SqlParsingAggregator(Closeable):
                     source=models.QuerySourceClass.SYSTEM,
                     created=query.make_created_audit_stamp(),
                     lastModified=query.make_last_modified_audit_stamp(),
+                    origin=query.origin.urn() if query.origin else None,
                 ),
                 models.QuerySubjectsClass(
                     subjects=[

datahub/sql_parsing/tool_meta_extractor.py

@@ -13,7 +13,7 @@ from datahub.api.entities.platformresource.platform_resource import (
 )
 from datahub.ingestion.api.report import Report
 from datahub.ingestion.graph.client import DataHubGraph
-from datahub.metadata.urns import CorpGroupUrn, CorpUserUrn
+from datahub.metadata.urns import CorpGroupUrn, CorpUserUrn, DataPlatformUrn, Urn
 from datahub.utilities.search_utils import LogicalOperator
 from datahub.utilities.stats_collections import int_top_k_dict
 
@@ -21,6 +21,10 @@ UrnStr = str
 
 logger = logging.getLogger(__name__)
 
+MODE_PLATFORM_URN = DataPlatformUrn.from_string("urn:li:dataPlatform:mode")
+LOOKER_PLATFORM_URN = DataPlatformUrn.from_string("urn:li:dataPlatform:looker")
+HEX_PLATFORM_URN = DataPlatformUrn.from_string("urn:li:dataPlatform:hex")
+
 
 class QueryLog(Protocol):
     """Represents Query Log Entry
@@ -30,6 +34,7 @@ class QueryLog(Protocol):
     query_text: str
     user: Optional[Union[CorpUserUrn, CorpGroupUrn]]
     extra_info: Optional[dict]
+    origin: Optional[Urn]
 
 
 def _get_last_line(query: str) -> str:
@@ -67,6 +72,10 @@ class ToolMetaExtractor:
                 "looker",
                 self._extract_looker_query,
             ),
+            (
+                "hex",
+                self._extract_hex_query,
+            ),
         ]
         # maps user id (as string) to email address
         self.looker_user_mapping = looker_user_mapping
@@ -153,7 +162,7 @@ class ToolMetaExtractor:
         entry.extra_info = entry.extra_info or {}
         entry.extra_info["user_via"] = original_user
 
-        # TODO: Generate an "origin" urn.
+        entry.origin = MODE_PLATFORM_URN
 
         return True
 
@@ -190,6 +199,22 @@ class ToolMetaExtractor:
         entry.extra_info = entry.extra_info or {}
         entry.extra_info["user_via"] = original_user
 
+        entry.origin = LOOKER_PLATFORM_URN
+
+        return True
+
+    def _extract_hex_query(self, entry: QueryLog) -> bool:
+        """
+        Returns:
+            bool: whether QueryLog entry is that of hex.
+        """
+        last_line = _get_last_line(entry.query_text)
+
+        if not last_line.startswith("-- Hex query metadata:"):
+            return False
+
+        entry.origin = HEX_PLATFORM_URN
+
         return True
 
     def extract_bi_metadata(self, entry: QueryLog) -> bool:

datahub/testing/mcp_diff.py

@@ -8,7 +8,6 @@ import deepdiff.serialization
 import yaml
 from deepdiff import DeepDiff
 from deepdiff.model import DiffLevel
-from deepdiff.operator import BaseOperator
 from typing_extensions import Literal
 
 ReportType = Literal[
@@ -59,27 +58,12 @@ class AspectForDiff:
 
 @dataclasses.dataclass
 class DeltaInfo:
-    """Information about an MCP used to construct a diff delta.
-
-    In a separate class so it can be ignored by DeepDiff via MCPDeltaInfoOperator.
-    """
+    """Information about an MCP used to construct a diff delta."""
 
     idx: int  # Location in list of MCEs in golden file
     original: Dict[str, Any]  # Original json-serialized MCP
 
 
-class DeltaInfoOperator(BaseOperator):
-    """Warning: Doesn't seem to be working right now.
-    Ignored via an ignore path as an extra layer of defense.
-    """
-
-    def __init__(self):
-        super().__init__(types=[DeltaInfo])
-
-    def give_up_diffing(self, *args: Any, **kwargs: Any) -> bool:
-        return True
-
-
 AspectsByUrn = Dict[str, Dict[str, List[AspectForDiff]]]
 
 
@@ -176,7 +160,6 @@ class MCPDiff:
                     t2=t2,
                     exclude_regex_paths=ignore_paths,
                     ignore_order=True,
-                    custom_operators=[DeltaInfoOperator()],
                 )
                 if diff:
                     aspect_changes[urn][aspect_name] = MCPAspectDiff.create(diff)