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

datahub/sql_parsing/sqlglot_utils.py

@@ -1,7 +1,6 @@
 from datahub.sql_parsing._sqlglot_patch import SQLGLOT_PATCHED
 
 import functools
-import hashlib
 import logging
 import re
 from typing import Dict, Iterable, Optional, Tuple, Union
@@ -10,6 +9,8 @@ import sqlglot
 import sqlglot.errors
 import sqlglot.optimizer.eliminate_ctes
 
+from datahub.sql_parsing.fingerprint_utils import generate_hash
+
 assert SQLGLOT_PATCHED
 
 logger = logging.getLogger(__name__)
@@ -39,9 +40,6 @@ def _get_dialect_str(platform: str) -> str:
         # let the fuzzy resolution logic handle it.
         # MariaDB is a fork of MySQL, so we reuse the same dialect.
         return "mysql, normalization_strategy = lowercase"
-    # Dremio is based upon drill. Not 100% compatibility
-    elif platform == "dremio":
-        return "drill"
     else:
         return platform
 
@@ -114,6 +112,8 @@ def _expression_to_string(
     return expression.sql(dialect=get_dialect(platform))
 
 
+PLACEHOLDER_BACKWARD_FINGERPRINT_NORMALIZATION = re.compile(r"(%s|\$\d|\?)")
+
 _BASIC_NORMALIZATION_RULES = {
     # Remove /* */ comments.
     re.compile(r"/\*.*?\*/", re.DOTALL): "",
@@ -129,7 +129,9 @@ _BASIC_NORMALIZATION_RULES = {
     re.compile(r"'[^']*'"): "?",
     # Replace sequences of IN/VALUES with a single placeholder.
     # The r" ?" makes it more robust to uneven spacing.
-    re.compile(r"\b(IN|VALUES)\s*\( ?\?(?:, ?\?)* ?\)", re.IGNORECASE): r"\1 (?)",
+    re.compile(
+        r"\b(IN|VALUES)\s*\( ?(?:%s|\$\d|\?)(?:, ?(?:%s|\$\d|\?))* ?\)", re.IGNORECASE
+    ): r"\1 (?)",
     # Normalize parenthesis spacing.
     re.compile(r"\( "): "(",
     re.compile(r" \)"): ")",
@@ -138,6 +140,9 @@ _BASIC_NORMALIZATION_RULES = {
     # e.g. "col1,col2" -> "col1, col2"
     re.compile(r"\b ,"): ",",
     re.compile(r"\b,\b"): ", ",
+    # MAKE SURE THAT THIS IS AFTER THE ABOVE REPLACEMENT
+    # Replace all versions of placeholders with generic ? placeholder.
+    PLACEHOLDER_BACKWARD_FINGERPRINT_NORMALIZATION: "?",
 }
 _TABLE_NAME_NORMALIZATION_RULES = {
     # Replace UUID-like strings with a placeholder (both - and _ variants).
@@ -251,18 +256,20 @@ def generalize_query(expression: sqlglot.exp.ExpOrStr, dialect: DialectOrStr) ->
     return expression.transform(_strip_expression, copy=True).sql(dialect=dialect)
 
 
-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()
-
-
 def get_query_fingerprint_debug(
-    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr, fast: bool = False
+    expression: sqlglot.exp.ExpOrStr,
+    platform: DialectOrStr,
+    fast: bool = False,
+    secondary_id: Optional[str] = None,
 ) -> Tuple[str, Optional[str]]:
     try:
         if not fast:
             dialect = get_dialect(platform)
             expression_sql = generalize_query(expression, dialect=dialect)
+            # Normalize placeholders for consistent fingerprinting -> this only needs to be backward compatible with earlier sqglot generated generalized queries where the placeholders were always ?
+            expression_sql = PLACEHOLDER_BACKWARD_FINGERPRINT_NORMALIZATION.sub(
+                "?", expression_sql
+            )
         else:
             expression_sql = generalize_query_fast(expression, dialect=platform)
     except (ValueError, sqlglot.errors.SqlglotError) as e:
@@ -272,16 +279,18 @@ def get_query_fingerprint_debug(
         logger.debug("Failed to generalize query for fingerprinting: %s", e)
         expression_sql = None
 
-    fingerprint = generate_hash(
-        expression_sql
-        if expression_sql is not None
-        else _expression_to_string(expression, platform=platform)
-    )
+    text = expression_sql or _expression_to_string(expression, platform=platform)
+    if secondary_id:
+        text = text + " -- " + secondary_id
+    fingerprint = generate_hash(text=text)
     return fingerprint, expression_sql
 
 
 def get_query_fingerprint(
-    expression: sqlglot.exp.ExpOrStr, platform: DialectOrStr, fast: bool = False
+    expression: sqlglot.exp.ExpOrStr,
+    platform: DialectOrStr,
+    fast: bool = False,
+    secondary_id: Optional[str] = None,
 ) -> str:
     """Get a fingerprint for a SQL query.
 
@@ -298,12 +307,15 @@ def get_query_fingerprint(
     Args:
         expression: The SQL query to fingerprint.
         platform: The SQL dialect to use.
+        secondary_id: An optional additional id string to included in the final fingerprint.
 
     Returns:
         The fingerprint for the SQL query.
     """
 
-    return get_query_fingerprint_debug(expression, platform, fast=fast)[0]
+    return get_query_fingerprint_debug(
+        expression=expression, platform=platform, fast=fast, secondary_id=secondary_id
+    )[0]
 
 
 @functools.lru_cache(maxsize=FORMAT_QUERY_CACHE_SIZE)

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,20 @@ 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.
+        """
+        if "-- Hex query metadata:" not in entry.query_text:
+            return False
+
+        entry.origin = HEX_PLATFORM_URN
+
         return True
 
     def extract_bi_metadata(self, entry: QueryLog) -> bool:

datahub/telemetry/telemetry.py

@@ -16,6 +16,11 @@ from datahub._version import __version__, nice_version_name
 from datahub.cli.config_utils import DATAHUB_ROOT_FOLDER
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import ExceptionWithProps
+from datahub.configuration.env_vars import (
+    get_sentry_dsn,
+    get_sentry_environment,
+    get_telemetry_timeout,
+)
 from datahub.metadata.schema_classes import _custom_package_path
 from datahub.utilities.perf_timer import PerfTimer
 
@@ -97,14 +102,14 @@ if any(var in os.environ for var in CI_ENV_VARS):
 if _custom_package_path:
     ENV_ENABLED = False
 
-TIMEOUT = int(os.environ.get("DATAHUB_TELEMETRY_TIMEOUT", "10"))
+TIMEOUT = int(get_telemetry_timeout())
 MIXPANEL_ENDPOINT = "track.datahubproject.io/mp"
 MIXPANEL_TOKEN = "5ee83d940754d63cacbf7d34daa6f44a"
-SENTRY_DSN: Optional[str] = os.environ.get("SENTRY_DSN", None)
-SENTRY_ENVIRONMENT: str = os.environ.get("SENTRY_ENVIRONMENT", "dev")
+SENTRY_DSN: Optional[str] = get_sentry_dsn()
+SENTRY_ENVIRONMENT: str = get_sentry_environment()
 
 
-def _default_telemetry_properties() -> Dict[str, Any]:
+def _default_global_properties() -> Dict[str, Any]:
     return {
         "datahub_version": nice_version_name(),
         "python_version": platform.python_version(),
@@ -122,6 +127,7 @@ class Telemetry:
     context_properties: Dict[str, Any] = {}
 
     def __init__(self):
+        self.global_properties = _default_global_properties()
         self.context_properties = {}
 
         if SENTRY_DSN:
@@ -247,6 +253,10 @@ class Telemetry:
 
         return False
 
+    def add_global_property(self, key: str, value: Any) -> None:
+        self.global_properties[key] = value
+        self._update_sentry_properties()
+
     def set_context(
         self,
         server: Optional["DataHubGraph"] = None,
@@ -257,16 +267,20 @@ class Telemetry:
             **(properties or {}),
         }
 
-        if self.sentry_enabled:
-            from sentry_sdk import set_tag
+        self._update_sentry_properties()
 
-            properties = {
-                **_default_telemetry_properties(),
-                **self.context_properties,
-            }
+    def _update_sentry_properties(self) -> None:
+        properties = {
+            **self.global_properties,
+            **self.context_properties,
+        }
+        if self.sentry_enabled:
+            import sentry_sdk
 
-            for key in properties:
-                set_tag(key, properties[key])
+            # Note: once we're on sentry-sdk 2.1.0+, we can use sentry_sdk.set_tags(properties)
+            # See https://github.com/getsentry/sentry-python/commit/6c960d752c7c7aff3fd7469d2e9ad98f19663aa8
+            for key, value in properties.items():
+                sentry_sdk.set_tag(key, value)
 
     def init_capture_exception(self) -> None:
         if self.sentry_enabled:
@@ -300,7 +314,7 @@ class Telemetry:
         try:
             self.mp.people_set(
                 self.client_id,
-                _default_telemetry_properties(),
+                self.global_properties,
             )
         except Exception as e:
             logger.debug(f"Error initializing telemetry: {e}")
@@ -334,7 +348,7 @@ class Telemetry:
                 logger.debug(f"Sending telemetry for {event_name}")
 
             properties = {
-                **_default_telemetry_properties(),
+                **self.global_properties,
                 **self.context_properties,
                 **properties,
             }
@@ -352,10 +366,10 @@ class Telemetry:
             }
         else:
             return {
-                "server_type": server.server_config.get("datahub", {}).get(
+                "server_type": server.server_config.raw_config.get("datahub", {}).get(
                     "serverType", "missing"
                 ),
-                "server_version": server.server_config.get("versions", {})
+                "server_version": server.server_config.raw_config.get("versions", {})
                 .get("acryldata/datahub", {})
                 .get("version", "missing"),
                 "server_id": server.server_id or "missing",

datahub/testing/check_imports.py

@@ -9,7 +9,7 @@ def ensure_no_indirect_model_imports(dirs: List[pathlib.Path]) -> None:
     # If our needs become more complex, we should move to a proper linter.
     denied_imports = {
         "src.": "datahub.*",
-        "datahub.metadata._schema_classes": "datahub.metadata.schema_classes",
+        "datahub.metadata._internal_schema_classes": "datahub.metadata.schema_classes",
         "datahub.metadata._urns": "datahub.metadata.urns",
     }
     ignored_files = {

datahub/testing/docker_utils.py

@@ -48,16 +48,22 @@ def wait_for_port(
         subprocess.run(f"docker logs {container_name}", shell=True, check=True)
 
 
+DOCKER_DEFAULT_UNLIMITED_PARALLELISM = -1
+
+
 @pytest.fixture(scope="module")
 def docker_compose_runner(
     docker_compose_command, docker_compose_project_name, docker_setup, docker_cleanup
 ):
     @contextlib.contextmanager
     def run(
-        compose_file_path: Union[str, List[str]], key: str, cleanup: bool = True
+        compose_file_path: Union[str, List[str]],
+        key: str,
+        cleanup: bool = True,
+        parallel: int = DOCKER_DEFAULT_UNLIMITED_PARALLELISM,
     ) -> Iterator[pytest_docker.plugin.Services]:
         with pytest_docker.plugin.get_docker_services(
-            docker_compose_command=docker_compose_command,
+            docker_compose_command=f"{docker_compose_command} --parallel {parallel}",
             # We can remove the type ignore once this is merged:
             # https://github.com/avast/pytest-docker/pull/108
             docker_compose_file=compose_file_path,  # type: ignore