acryl-datahub 0.15.0rc25__py3-none-any.whl → 0.15.0.1__py3-none-any.whl

datahub/specific/structured_property.py

@@ -29,7 +29,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
         self._add_patch(
             StructuredPropertyDefinition.ASPECT_NAME,
             "add",
-            path="/qualifiedName",
+            path=("qualifiedName",),
             value=qualified_name,
         )
         return self
@@ -41,7 +41,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
             self._add_patch(
                 StructuredPropertyDefinition.ASPECT_NAME,
                 "add",
-                path="/displayName",
+                path=("displayName",),
                 value=display_name,
             )
         return self
@@ -53,7 +53,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
         self._add_patch(
             StructuredPropertyDefinition.ASPECT_NAME,
             "add",
-            path="/valueType",
+            path=("valueType",),
             value=value_type,
         )
         return self
@@ -66,7 +66,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
             self._add_patch(
                 StructuredPropertyDefinition.ASPECT_NAME,
                 "add",
-                path="/typeQualifier",
+                path=("typeQualifier",),
                 value=type_qualifier,
             )
         return self
@@ -78,7 +78,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
         self._add_patch(
             StructuredPropertyDefinition.ASPECT_NAME,
             "add",
-            path=f"/allowedValues/{str(allowed_value.get('value'))}",
+            path=("allowedValues", str(allowed_value.get("value"))),
             value=allowed_value,
         )
         return self
@@ -87,7 +87,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
         self._add_patch(
             StructuredPropertyDefinition.ASPECT_NAME,
             "add",
-            path="/cardinality",
+            path=("cardinality",),
             value=cardinality,
         )
         return self
@@ -98,7 +98,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
         self._add_patch(
             StructuredPropertyDefinition.ASPECT_NAME,
             "add",
-            path=f"/entityTypes/{self.quote(str(entity_type))}",
+            path=("entityTypes", str(entity_type)),
             value=entity_type,
         )
         return self
@@ -110,7 +110,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
             self._add_patch(
                 StructuredPropertyDefinition.ASPECT_NAME,
                 "add",
-                path="/description",
+                path=("description",),
                 value=description,
             )
         return self
@@ -119,7 +119,7 @@ class StructuredPropertyPatchBuilder(MetadataPatchProposal):
         self._add_patch(
             StructuredPropertyDefinition.ASPECT_NAME,
             "add",
-            path="/immutable",
+            path=("immutable",),
             value=immutable,
         )
         return self

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -165,6 +165,7 @@ class KnownQueryLineageInfo:
     timestamp: Optional[datetime] = None
     session_id: Optional[str] = None
     query_type: QueryType = QueryType.UNKNOWN
+    query_id: Optional[str] = None
 
 
 @dataclasses.dataclass
@@ -198,7 +199,7 @@ class TableSwap:
 
 @dataclasses.dataclass
 class PreparsedQuery:
-    # If not provided, we will generate one using the fast fingerprint generator.
+    # If not provided, we will generate one using the fingerprint generator.
     query_id: Optional[QueryId]
 
     query_text: str
@@ -490,7 +491,7 @@ class SqlParsingAggregator(Closeable):
             self._exit_stack.push(self._query_usage_counts)
 
         # Tool Extractor
-        self._tool_meta_extractor = ToolMetaExtractor()
+        self._tool_meta_extractor = ToolMetaExtractor.create(graph)
         self.report.tool_meta_report = self._tool_meta_extractor.report
 
     def close(self) -> None:
@@ -618,12 +619,13 @@ class SqlParsingAggregator(Closeable):
         self.report.num_known_query_lineage += 1
 
         # Generate a fingerprint for the query.
-        with self.report.sql_fingerprinting_timer:
-            query_fingerprint = get_query_fingerprint(
-                known_query_lineage.query_text,
-                platform=self.platform.platform_name,
-                fast=True,
-            )
+        query_fingerprint = known_query_lineage.query_id
+        if not query_fingerprint:
+            with self.report.sql_fingerprinting_timer:
+                query_fingerprint = get_query_fingerprint(
+                    known_query_lineage.query_text,
+                    platform=self.platform.platform_name,
+                )
         formatted_query = self._maybe_format_query(known_query_lineage.query_text)
 
         # Register the query.
@@ -848,7 +850,6 @@ class SqlParsingAggregator(Closeable):
             query_fingerprint = get_query_fingerprint(
                 parsed.query_text,
                 platform=self.platform.platform_name,
-                fast=True,
             )
 
         # Format the query.

datahub/sql_parsing/sqlglot_lineage.py

@@ -66,6 +66,7 @@ SQL_LINEAGE_TIMEOUT_ENABLED = get_boolean_env_variable(
     "SQL_LINEAGE_TIMEOUT_ENABLED", True
 )
 SQL_LINEAGE_TIMEOUT_SECONDS = 10
+SQL_PARSER_TRACE = get_boolean_env_variable("DATAHUB_SQL_PARSER_TRACE", False)
 
 
 # These rules are a subset of the rules in sqlglot.optimizer.optimizer.RULES.
@@ -365,10 +366,11 @@ def _prepare_query_columns(
 
             return node
 
-        # logger.debug(
-        #     "Prior to case normalization sql %s",
-        #     statement.sql(pretty=True, dialect=dialect),
-        # )
+        if SQL_PARSER_TRACE:
+            logger.debug(
+                "Prior to case normalization sql %s",
+                statement.sql(pretty=True, dialect=dialect),
+            )
         statement = statement.transform(_sqlglot_force_column_normalizer, copy=False)
         # logger.debug(
         #     "Sql after casing normalization %s",
@@ -562,7 +564,7 @@ def _select_statement_cll(  # noqa: C901
                 )
             )
 
-        # TODO: Also extract referenced columns (aka auxillary / non-SELECT lineage)
+        # TODO: Also extract referenced columns (aka auxiliary / non-SELECT lineage)
     except (sqlglot.errors.OptimizeError, ValueError, IndexError) as e:
         raise SqlUnderstandingError(
             f"sqlglot failed to compute some lineage: {e}"
@@ -1022,6 +1024,14 @@ def _sqlglot_lineage_inner(
     logger.debug(
         f"Resolved {total_schemas_resolved} of {total_tables_discovered} table schemas"
     )
+    if SQL_PARSER_TRACE:
+        for qualified_table, schema_info in table_name_schema_mapping.items():
+            logger.debug(
+                "Table name %s resolved to %s with schema %s",
+                qualified_table,
+                table_name_urn_mapping[qualified_table],
+                schema_info,
+            )
 
     column_lineage: Optional[List[_ColumnLineageInfo]] = None
     try:

datahub/sql_parsing/tool_meta_extractor.py

@@ -1,3 +1,4 @@
+import contextlib
 import json
 import logging
 from dataclasses import dataclass, field
@@ -5,8 +6,15 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 
 from typing_extensions import Protocol
 
+from datahub.api.entities.platformresource.platform_resource import (
+    ElasticPlatformResourceQuery,
+    PlatformResource,
+    PlatformResourceSearchFields,
+)
 from datahub.ingestion.api.report import Report
+from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.urns import CorpGroupUrn, CorpUserUrn
+from datahub.utilities.search_utils import LogicalOperator
 from datahub.utilities.stats_collections import int_top_k_dict
 
 UrnStr = str
@@ -31,6 +39,8 @@ def _get_last_line(query: str) -> str:
 @dataclass
 class ToolMetaExtractorReport(Report):
     num_queries_meta_extracted: Dict[str, int] = field(default_factory=int_top_k_dict)
+    failures: List[str] = field(default_factory=list)
+    looker_user_mapping_missing: Optional[bool] = None
 
 
 class ToolMetaExtractor:
@@ -42,14 +52,83 @@ class ToolMetaExtractor:
     by warehouse query logs.
     """
 
-    def __init__(self) -> None:
-        self.report = ToolMetaExtractorReport()
+    def __init__(
+        self,
+        report: ToolMetaExtractorReport,
+        looker_user_mapping: Optional[Dict[str, str]] = None,
+    ) -> None:
+        self.report = report
         self.known_tool_extractors: List[Tuple[str, Callable[[QueryLog], bool]]] = [
             (
                 "mode",
                 self._extract_mode_query,
-            )
+            ),
+            (
+                "looker",
+                self._extract_looker_query,
+            ),
         ]
+        # maps user id (as string) to email address
+        self.looker_user_mapping = looker_user_mapping
+
+    @classmethod
+    def create(
+        cls,
+        graph: Optional[DataHubGraph] = None,
+    ) -> "ToolMetaExtractor":
+        report = ToolMetaExtractorReport()
+        looker_user_mapping = None
+        if graph:
+            try:
+                looker_user_mapping = cls.extract_looker_user_mapping_from_graph(
+                    graph, report
+                )
+            except Exception as e:
+                report.failures.append(
+                    f"Unexpected error during Looker user metadata extraction: {str(e)}"
+                )
+
+        return cls(report, looker_user_mapping)
+
+    @classmethod
+    def extract_looker_user_mapping_from_graph(
+        cls, graph: DataHubGraph, report: ToolMetaExtractorReport
+    ) -> Optional[Dict[str, str]]:
+        looker_user_mapping = None
+        query = (
+            ElasticPlatformResourceQuery.create_from()
+            .group(LogicalOperator.AND)
+            .add_field_match(PlatformResourceSearchFields.PLATFORM, "looker")
+            .add_field_match(
+                PlatformResourceSearchFields.RESOURCE_TYPE,
+                "USER_ID_MAPPING",
+            )
+            .end()
+        )
+        platform_resources = list(
+            PlatformResource.search_by_filters(query=query, graph_client=graph)
+        )
+
+        if len(platform_resources) == 0:
+            report.looker_user_mapping_missing = True
+        elif len(platform_resources) > 1:
+            report.failures.append(
+                "Looker user metadata extraction failed. Found more than one looker user id mappings."
+            )
+        else:
+            platform_resource = platform_resources[0]
+
+            if (
+                platform_resource
+                and platform_resource.resource_info
+                and platform_resource.resource_info.value
+            ):
+                with contextlib.suppress(ValueError, AssertionError):
+                    value = platform_resource.resource_info.value.as_raw_json()
+                    if value:
+                        looker_user_mapping = value
+
+        return looker_user_mapping
 
     def _extract_mode_query(self, entry: QueryLog) -> bool:
         """
@@ -78,14 +157,49 @@ class ToolMetaExtractor:
 
         return True
 
+    def _extract_looker_query(self, entry: QueryLog) -> bool:
+        """
+        Returns:
+            bool: whether QueryLog entry is that of looker and looker user info
+            is extracted into entry.
+        """
+        if not self.looker_user_mapping:
+            return False
+
+        last_line = _get_last_line(entry.query_text)
+
+        if not (last_line.startswith("--") and "Looker Query Context" in last_line):
+            return False
+
+        start_quote_idx = last_line.index("'")
+        end_quote_idx = last_line.rindex("'")
+        if start_quote_idx == -1 or end_quote_idx == -1:
+            return False
+
+        looker_json_raw = last_line[start_quote_idx + 1 : end_quote_idx]
+        looker_json = json.loads(looker_json_raw)
+
+        user_id = str(looker_json["user_id"])
+        email = self.looker_user_mapping.get(user_id)
+        if not email:
+            return False
+
+        original_user = entry.user
+
+        entry.user = email_to_user_urn(email)
+        entry.extra_info = entry.extra_info or {}
+        entry.extra_info["user_via"] = original_user
+
+        return True
+
     def extract_bi_metadata(self, entry: QueryLog) -> bool:
         for tool, meta_extractor in self.known_tool_extractors:
             try:
                 if meta_extractor(entry):
                     self.report.num_queries_meta_extracted[tool] += 1
                     return True
-            except Exception:
-                logger.debug("Tool metadata extraction failed with error : {e}")
+            except Exception as e:
+                logger.debug(f"Tool metadata extraction failed with error : {e}")
         return False
 
 

datahub/utilities/time.py

@@ -1,6 +1,8 @@
 import time
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime
+
+from datahub.emitter.mce_builder import make_ts_millis, parse_ts_millis
 
 
 def get_current_time_in_seconds() -> int:
@@ -9,12 +11,15 @@ def get_current_time_in_seconds() -> int:
 
 def ts_millis_to_datetime(ts_millis: int) -> datetime:
     """Converts input timestamp in milliseconds to a datetime object with UTC timezone"""
-    return datetime.fromtimestamp(ts_millis / 1000, tz=timezone.utc)
+    return parse_ts_millis(ts_millis)
 
 
 def datetime_to_ts_millis(dt: datetime) -> int:
     """Converts a datetime object to timestamp in milliseconds"""
-    return int(round(dt.timestamp() * 1000))
+    # TODO: Deprecate these helpers in favor of make_ts_millis and parse_ts_millis.
+    # The other ones support None with a typing overload.
+    # Also possibly move those helpers to this file.
+    return make_ts_millis(dt)
 
 
 @dataclass

datahub/utilities/urns/_urn_base.py

@@ -1,9 +1,10 @@
 import functools
 import urllib.parse
 from abc import abstractmethod
-from typing import ClassVar, Dict, List, Optional, Type, TypeVar
+from typing import ClassVar, Dict, List, Optional, Type
 
 from deprecated import deprecated
+from typing_extensions import Self
 
 from datahub.utilities.urns.error import InvalidUrnError
 
@@ -42,9 +43,6 @@ def _split_entity_id(entity_id: str) -> List[str]:
     return parts
 
 
-_UrnSelf = TypeVar("_UrnSelf", bound="Urn")
-
-
 @functools.total_ordering
 class Urn:
     """
@@ -88,7 +86,7 @@ class Urn:
         return self._entity_ids
 
     @classmethod
-    def from_string(cls: Type[_UrnSelf], urn_str: str) -> "_UrnSelf":
+    def from_string(cls, urn_str: str) -> Self:
         """
         Creates an Urn from its string representation.
 
@@ -174,7 +172,7 @@ class Urn:
 
     @classmethod
     @deprecated(reason="prefer .from_string")
-    def create_from_string(cls: Type[_UrnSelf], urn_str: str) -> "_UrnSelf":
+    def create_from_string(cls, urn_str: str) -> Self:
         return cls.from_string(urn_str)
 
     @deprecated(reason="prefer .entity_ids")
@@ -270,5 +268,5 @@ class _SpecificUrn(Urn):
 
     @classmethod
     @abstractmethod
-    def _parse_ids(cls: Type[_UrnSelf], entity_ids: List[str]) -> _UrnSelf:
+    def _parse_ids(cls, entity_ids: List[str]) -> Self:
         raise NotImplementedError()