acryl-datahub 1.2.0.2rc1__py3-none-any.whl → 1.2.0.2rc3__py3-none-any.whl

datahub/ingestion/source/unity/proxy.py

@@ -521,9 +521,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_schema_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching schema tags for catalog: {catalog}")
+        logger.info(f"Fetching schema tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.schema_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.schema_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -544,9 +544,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_catalog_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching table tags for catalog: {catalog}")
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.catalog_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.catalog_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -566,9 +566,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_table_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching table tags for catalog: {catalog}")
+        logger.info(f"Fetching table tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.table_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.table_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}
@@ -589,9 +589,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
     @cached(cachetools.FIFOCache(maxsize=100))
     def get_column_tags(self, catalog: str) -> Dict[str, List[UnityCatalogTag]]:
         """Optimized version using databricks-sql"""
-        logger.info(f"Fetching column tags for catalog: {catalog}")
+        logger.info(f"Fetching column tags for catalog: `{catalog}`")
 
-        query = f"SELECT * FROM {catalog}.information_schema.column_tags"
+        query = f"SELECT * FROM `{catalog}`.information_schema.column_tags"
         rows = self._execute_sql_query(query)
 
         result_dict: Dict[str, List[UnityCatalogTag]] = {}

datahub/metadata/_internal_schema_classes.py

@@ -20163,23 +20163,24 @@ class DataHubPageModuleVisibilityClass(DictWrapper):
     
     
 class HierarchyModuleParamsClass(DictWrapper):
-    """The params required if the module is type HIERARCHY_VIEW
-    TODO: add filters
-    relatedEntitiesFilter: optional Filter"""
+    """The params required if the module is type HIERARCHY_VIEW"""
     
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.module.HierarchyModuleParams")
     def __init__(self,
         showRelatedEntities: bool,
         assetUrns: Union[None, List[str]]=None,
+        relatedEntitiesFilterJson: Union[None, str]=None,
     ):
         super().__init__()
         
         self.assetUrns = assetUrns
         self.showRelatedEntities = showRelatedEntities
+        self.relatedEntitiesFilterJson = relatedEntitiesFilterJson
     
     def _restore_defaults(self) -> None:
         self.assetUrns = self.RECORD_SCHEMA.fields_dict["assetUrns"].default
         self.showRelatedEntities = bool()
+        self.relatedEntitiesFilterJson = self.RECORD_SCHEMA.fields_dict["relatedEntitiesFilterJson"].default
     
     
     @property
@@ -20202,6 +20203,20 @@ class HierarchyModuleParamsClass(DictWrapper):
         self._inner_dict['showRelatedEntities'] = value
     
     
+    @property
+    def relatedEntitiesFilterJson(self) -> Union[None, str]:
+        """Optional filters to filter relatedEntities (assetUrns) out
+    
+    The stringified json representing the logical predicate built in the UI to select assets.
+    This predicate is turned into orFilters to send through graphql since graphql doesn't support
+    arbitrary nesting. This string is used to restore the UI for this logical predicate."""
+        return self._inner_dict.get('relatedEntitiesFilterJson')  # type: ignore
+    
+    @relatedEntitiesFilterJson.setter
+    def relatedEntitiesFilterJson(self, value: Union[None, str]) -> None:
+        self._inner_dict['relatedEntitiesFilterJson'] = value
+    
+    
 class LinkModuleParamsClass(DictWrapper):
     # No docs available.
     

datahub/metadata/schema.avsc

@@ -17844,9 +17844,18 @@
                     {
                       "type": "boolean",
                       "name": "showRelatedEntities"
+                    },
+                    {
+                      "type": [
+                        "null",
+                        "string"
+                      ],
+                      "name": "relatedEntitiesFilterJson",
+                      "default": null,
+                      "doc": "Optional filters to filter relatedEntities (assetUrns) out\n\nThe stringified json representing the logical predicate built in the UI to select assets.\nThis predicate is turned into orFilters to send through graphql since graphql doesn't support\narbitrary nesting. This string is used to restore the UI for this logical predicate."
                     }
                   ],
-                  "doc": "The params required if the module is type HIERARCHY_VIEW\nTODO: add filters\nrelatedEntitiesFilter: optional Filter"
+                  "doc": "The params required if the module is type HIERARCHY_VIEW"
                 }
               ],
               "name": "hierarchyViewParams",

datahub/metadata/schemas/DataHubPageModuleProperties.avsc

@@ -181,9 +181,18 @@
                   {
                     "type": "boolean",
                     "name": "showRelatedEntities"
+                  },
+                  {
+                    "type": [
+                      "null",
+                      "string"
+                    ],
+                    "name": "relatedEntitiesFilterJson",
+                    "default": null,
+                    "doc": "Optional filters to filter relatedEntities (assetUrns) out\n\nThe stringified json representing the logical predicate built in the UI to select assets.\nThis predicate is turned into orFilters to send through graphql since graphql doesn't support\narbitrary nesting. This string is used to restore the UI for this logical predicate."
                   }
                 ],
-                "doc": "The params required if the module is type HIERARCHY_VIEW\nTODO: add filters\nrelatedEntitiesFilter: optional Filter"
+                "doc": "The params required if the module is type HIERARCHY_VIEW"
               }
             ],
             "name": "hierarchyViewParams",

datahub/sdk/dataset.py

@@ -72,6 +72,11 @@ UpstreamLineageInputType: TypeAlias = Union[
     Dict[DatasetUrnOrStr, ColumnLineageMapping],
 ]
 
+ViewDefinitionInputType: TypeAlias = Union[
+    str,
+    models.ViewPropertiesClass,
+]
+
 
 def _parse_upstream_input(
     upstream_input: UpstreamInputType,
@@ -467,6 +472,7 @@ class Dataset(
         custom_properties: Optional[Dict[str, str]] = None,
         created: Optional[datetime] = None,
         last_modified: Optional[datetime] = None,
+        view_definition: Optional[ViewDefinitionInputType] = None,
         # Standard aspects.
         parent_container: ParentContainerInputType | Unset = unset,
         subtype: Optional[str] = None,
@@ -495,6 +501,7 @@ class Dataset(
             custom_properties: Optional dictionary of custom properties.
             created: Optional creation timestamp.
             last_modified: Optional last modification timestamp.
+            view_definition: Optional view definition for the dataset.
             parent_container: Optional parent container for this dataset.
             subtype: Optional subtype of the dataset.
             owners: Optional list of owners.
@@ -536,6 +543,8 @@ class Dataset(
             self.set_created(created)
         if last_modified is not None:
             self.set_last_modified(last_modified)
+        if view_definition is not None:
+            self.set_view_definition(view_definition)
 
         if parent_container is not unset:
             self._set_container(parent_container)
@@ -717,6 +726,41 @@ class Dataset(
     def set_last_modified(self, last_modified: datetime) -> None:
         self._ensure_dataset_props().lastModified = make_time_stamp(last_modified)
 
+    @property
+    def view_definition(self) -> Optional[models.ViewPropertiesClass]:
+        """Get the view definition of the dataset.
+
+        Under typical usage, this will be present if the subtype is "View".
+
+        Returns:
+            The view definition if set, None otherwise.
+        """
+        return self._get_aspect(models.ViewPropertiesClass)
+
+    def set_view_definition(self, view_definition: ViewDefinitionInputType) -> None:
+        """Set the view definition of the dataset.
+
+        If you're setting a view definition, subtype should typically be set to "view".
+
+        If a string is provided, it will be treated as a SQL view definition. To set
+        a custom language or other properties, provide a ViewPropertiesClass object.
+
+        Args:
+            view_definition: The view definition to set.
+        """
+        if isinstance(view_definition, models.ViewPropertiesClass):
+            self._set_aspect(view_definition)
+        elif isinstance(view_definition, str):
+            self._set_aspect(
+                models.ViewPropertiesClass(
+                    materialized=False,
+                    viewLogic=view_definition,
+                    viewLanguage="SQL",
+                )
+            )
+        else:
+            assert_never(view_definition)
+
     def _schema_dict(self) -> Dict[str, models.SchemaFieldClass]:
         schema_metadata = self._get_aspect(models.SchemaMetadataClass)
         if schema_metadata is None:

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: