acryl-datahub 1.2.0.2rc2__py3-none-any.whl → 1.2.0.3__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import abc
+import json
 from typing import (
     TYPE_CHECKING,
     Annotated,
@@ -406,26 +407,45 @@ if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
 else:
     from pydantic import Discriminator, Tag
 
+    def _parse_json_from_string(value: Any) -> Any:
+        if isinstance(value, str):
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                return value
+        else:
+            return value
+
     # TODO: Once we're fully on pydantic 2, we can use a RootModel here.
     # That way we'd be able to attach methods to the Filter type.
     # e.g. replace load_filters(...) with Filter.load(...)
     Filter = Annotated[
-        Union[
-            Annotated[_And, Tag(_And._field_discriminator())],
-            Annotated[_Or, Tag(_Or._field_discriminator())],
-            Annotated[_Not, Tag(_Not._field_discriminator())],
-            Annotated[_EntityTypeFilter, Tag(_EntityTypeFilter._field_discriminator())],
-            Annotated[
-                _EntitySubtypeFilter, Tag(_EntitySubtypeFilter._field_discriminator())
+        Annotated[
+            Union[
+                Annotated[_And, Tag(_And._field_discriminator())],
+                Annotated[_Or, Tag(_Or._field_discriminator())],
+                Annotated[_Not, Tag(_Not._field_discriminator())],
+                Annotated[
+                    _EntityTypeFilter, Tag(_EntityTypeFilter._field_discriminator())
+                ],
+                Annotated[
+                    _EntitySubtypeFilter,
+                    Tag(_EntitySubtypeFilter._field_discriminator()),
+                ],
+                Annotated[_StatusFilter, Tag(_StatusFilter._field_discriminator())],
+                Annotated[_PlatformFilter, Tag(_PlatformFilter._field_discriminator())],
+                Annotated[_DomainFilter, Tag(_DomainFilter._field_discriminator())],
+                Annotated[
+                    _ContainerFilter, Tag(_ContainerFilter._field_discriminator())
+                ],
+                Annotated[_EnvFilter, Tag(_EnvFilter._field_discriminator())],
+                Annotated[
+                    _CustomCondition, Tag(_CustomCondition._field_discriminator())
+                ],
             ],
-            Annotated[_StatusFilter, Tag(_StatusFilter._field_discriminator())],
-            Annotated[_PlatformFilter, Tag(_PlatformFilter._field_discriminator())],
-            Annotated[_DomainFilter, Tag(_DomainFilter._field_discriminator())],
-            Annotated[_ContainerFilter, Tag(_ContainerFilter._field_discriminator())],
-            Annotated[_EnvFilter, Tag(_EnvFilter._field_discriminator())],
-            Annotated[_CustomCondition, Tag(_CustomCondition._field_discriminator())],
+            Discriminator(_filter_discriminator),
         ],
-        Discriminator(_filter_discriminator),
+        pydantic.BeforeValidator(_parse_json_from_string),
     ]
 
     # Required to resolve forward references to "Filter"

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -49,6 +49,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,
@@ -109,6 +110,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
@@ -834,6 +836,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(
@@ -1168,6 +1171,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(
@@ -1175,6 +1179,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
 

datahub/telemetry/telemetry.py

@@ -272,7 +272,10 @@ class Telemetry:
         if self.sentry_enabled:
             import sentry_sdk
 
-            sentry_sdk.set_tags(properties)
+            # 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: