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

datahub/ingestion/source/unity/tag_entities.py

@@ -0,0 +1,197 @@
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+if TYPE_CHECKING:
+    from datahub.ingestion.source.unity.platform_resource_repository import (
+        UnityCatalogPlatformResourceRepository,
+    )
+
+from pydantic import BaseModel
+
+from datahub.api.entities.external.external_entities import (
+    ExternalEntity,
+    ExternalEntityId,
+    LinkedResourceSet,
+)
+from datahub.api.entities.external.unity_catalog_external_entites import UnityCatalogTag
+from datahub.api.entities.platformresource.platform_resource import (
+    PlatformResource,
+    PlatformResourceKey,
+)
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.urns import TagUrn
+from datahub.utilities.urns.urn import Urn
+
+
+class UnityCatalogTagSyncContext(BaseModel):
+    # it is intentionally empty
+    platform_instance: Optional[str] = None
+
+
+logger = logging.getLogger(__name__)
+
+
+class UnityCatalogTagPlatformResourceId(ExternalEntityId):
+    """
+    A Unity Catalog tag platform resource ID.
+    """
+
+    tag_key: str
+    tag_value: Optional[str] = None
+    platform_instance: Optional[str] = None
+    exists_in_unity_catalog: bool = False
+    persisted: bool = False
+
+    # this is a hack to make sure the property is a string and not private pydantic field
+    @staticmethod
+    def _RESOURCE_TYPE() -> str:
+        return "UnityCatalogTagPlatformResource"
+
+    def to_platform_resource_key(self) -> PlatformResourceKey:
+        return PlatformResourceKey(
+            platform="databricks",
+            resource_type=str(UnityCatalogTagPlatformResourceId._RESOURCE_TYPE()),
+            primary_key=f"{self.tag_key}:{self.tag_value}",
+            platform_instance=self.platform_instance,
+        )
+
+    @classmethod
+    def get_or_create_from_tag(
+        cls,
+        tag: UnityCatalogTag,
+        platform_resource_repository: "UnityCatalogPlatformResourceRepository",
+        exists_in_unity_catalog: bool = False,
+    ) -> "UnityCatalogTagPlatformResourceId":
+        """
+        Creates a UnityCatalogTagPlatformResourceId from a UnityCatalogTag.
+        """
+
+        existing_platform_resource = platform_resource_repository.search_entity_by_urn(
+            tag.to_datahub_tag_urn().urn()
+        )
+        if existing_platform_resource:
+            logger.debug(
+                f"Found existing UnityCatalogTagPlatformResourceId for tag {tag.key.raw_text}: {existing_platform_resource}"
+            )
+            return existing_platform_resource
+
+        return UnityCatalogTagPlatformResourceId(
+            tag_key=tag.key.raw_text,
+            tag_value=tag.value.raw_text if tag.value is not None else None,
+            platform_instance=platform_resource_repository.platform_instance,
+            exists_in_unity_catalog=exists_in_unity_catalog,
+            persisted=False,
+        )
+
+    @classmethod
+    def from_datahub_urn(
+        cls,
+        urn: str,
+        tag_sync_context: UnityCatalogTagSyncContext,
+        platform_resource_repository: "UnityCatalogPlatformResourceRepository",
+        graph: DataHubGraph,
+    ) -> "UnityCatalogTagPlatformResourceId":
+        """
+        Creates a UnityCatalogTagPlatformResourceId from a DataHub URN.
+        """
+        existing_platform_resource_id = (
+            platform_resource_repository.search_entity_by_urn(urn)
+        )
+        if existing_platform_resource_id:
+            return existing_platform_resource_id
+
+        new_unity_catalog_tag_id = cls.generate_tag_id(graph, tag_sync_context, urn)
+        if new_unity_catalog_tag_id:
+            resource_key = platform_resource_repository.get(
+                new_unity_catalog_tag_id.to_platform_resource_key()
+            )
+            if resource_key:
+                # Create a new ID with the correct state instead of mutating
+                return UnityCatalogTagPlatformResourceId(
+                    tag_key=new_unity_catalog_tag_id.tag_key,
+                    tag_value=new_unity_catalog_tag_id.tag_value,
+                    platform_instance=new_unity_catalog_tag_id.platform_instance,
+                    exists_in_unity_catalog=True,  # This tag exists in Unity Catalog
+                    persisted=new_unity_catalog_tag_id.persisted,
+                )
+            return new_unity_catalog_tag_id
+        raise ValueError(
+            f"Unable to create Unity Catalog tag ID from DataHub URN: {urn}"
+        )
+
+    @classmethod
+    def generate_tag_id(
+        cls, graph: DataHubGraph, tag_sync_context: UnityCatalogTagSyncContext, urn: str
+    ) -> "UnityCatalogTagPlatformResourceId":
+        parsed_urn = Urn.from_string(urn)
+        entity_type = parsed_urn.entity_type
+        if entity_type == "tag":
+            return UnityCatalogTagPlatformResourceId.from_datahub_tag(
+                TagUrn.from_string(urn), tag_sync_context
+            )
+        else:
+            raise ValueError(f"Unsupported entity type {entity_type} for URN {urn}")
+
+    @classmethod
+    def from_datahub_tag(
+        cls, tag_urn: TagUrn, tag_sync_context: UnityCatalogTagSyncContext
+    ) -> "UnityCatalogTagPlatformResourceId":
+        uc_tag = UnityCatalogTag.from_urn(tag_urn)
+
+        return UnityCatalogTagPlatformResourceId(
+            tag_key=str(uc_tag.key),
+            tag_value=str(uc_tag.value) if uc_tag.value is not None else None,
+            platform_instance=tag_sync_context.platform_instance,
+            exists_in_unity_catalog=False,
+        )
+
+
+class UnityCatalogTagPlatformResource(ExternalEntity):
+    datahub_urns: LinkedResourceSet
+    managed_by_datahub: bool
+    id: UnityCatalogTagPlatformResourceId
+    allowed_values: Optional[List[str]] = None
+
+    def get_id(self) -> ExternalEntityId:
+        return self.id
+
+    def is_managed_by_datahub(self) -> bool:
+        return self.managed_by_datahub
+
+    def datahub_linked_resources(self) -> LinkedResourceSet:
+        return self.datahub_urns
+
+    def as_platform_resource(self) -> PlatformResource:
+        return PlatformResource.create(
+            key=self.id.to_platform_resource_key(),
+            secondary_keys=[u for u in self.datahub_urns.urns],
+            value=self,
+        )
+
+    @classmethod
+    def create_default(
+        cls,
+        entity_id: ExternalEntityId,
+        managed_by_datahub: bool,
+    ) -> "UnityCatalogTagPlatformResource":
+        """Create a default Unity Catalog tag entity when none found in DataHub."""
+        # Type narrowing: we know this will be a UnityCatalogTagPlatformResourceId
+        assert isinstance(entity_id, UnityCatalogTagPlatformResourceId), (
+            f"Expected UnityCatalogTagPlatformResourceId, got {type(entity_id)}"
+        )
+
+        # Create a new entity ID with correct default state instead of mutating
+        default_entity_id = UnityCatalogTagPlatformResourceId(
+            tag_key=entity_id.tag_key,
+            tag_value=entity_id.tag_value,
+            platform_instance=entity_id.platform_instance,
+            exists_in_unity_catalog=False,  # New entities don't exist in Unity Catalog yet
+            persisted=False,  # New entities are not persisted yet
+        )
+
+        return cls(
+            id=default_entity_id,
+            datahub_urns=LinkedResourceSet(urns=[]),
+            managed_by_datahub=managed_by_datahub,
+            allowed_values=None,
+        )

datahub/ingestion/source/unity/usage.py

@@ -11,7 +11,10 @@ from databricks.sdk.service.sql import QueryStatementType
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.source_helpers import auto_empty_dataset_usage_statistics
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.unity.config import UnityCatalogSourceConfig
+from datahub.ingestion.source.unity.config import (
+    UnityCatalogSourceConfig,
+    UsageDataSource,
+)
 from datahub.ingestion.source.unity.proxy import UnityCatalogApiProxy
 from datahub.ingestion.source.unity.proxy_types import (
     OPERATION_STATEMENT_TYPES,
@@ -164,11 +167,50 @@ class UnityCatalogUsageExtractor:
                     aspect=operation_aspect,
                 ).as_workunit()
 
+    def _validate_usage_data_source_config(self) -> None:
+        """Validate usage data source configuration before execution."""
+        usage_data_source = self.config.usage_data_source
+
+        if (
+            usage_data_source == UsageDataSource.SYSTEM_TABLES
+            and not self.proxy.warehouse_id
+        ):
+            raise ValueError(
+                "usage_data_source is set to SYSTEM_TABLES but warehouse_id is not configured. "
+                "Either set warehouse_id or use AUTO/API mode."
+            )
+
     def _get_queries(self) -> Iterable[Query]:
         try:
-            yield from self.proxy.query_history(
-                self.config.start_time, self.config.end_time
-            )
+            self._validate_usage_data_source_config()
+            usage_data_source = self.config.usage_data_source
+
+            if usage_data_source == UsageDataSource.AUTO:
+                if self.proxy.warehouse_id:
+                    logger.info(
+                        "Using system tables for usage query history (AUTO mode)"
+                    )
+                    yield from self.proxy.get_query_history_via_system_tables(
+                        self.config.start_time, self.config.end_time
+                    )
+                else:
+                    logger.info(
+                        "Using API for usage query history (AUTO mode, no warehouse)"
+                    )
+                    yield from self.proxy.query_history(
+                        self.config.start_time, self.config.end_time
+                    )
+            elif usage_data_source == UsageDataSource.SYSTEM_TABLES:
+                logger.info("Using system tables for usage query history (forced)")
+                yield from self.proxy.get_query_history_via_system_tables(
+                    self.config.start_time, self.config.end_time
+                )
+            elif usage_data_source == UsageDataSource.API:
+                logger.info("Using API for usage query history (forced)")
+                yield from self.proxy.query_history(
+                    self.config.start_time, self.config.end_time
+                )
+
         except Exception as e:
             logger.warning("Error getting queries", exc_info=True)
             self.report.report_warning("get-queries", str(e))

datahub/ingestion/source/usage/clickhouse_usage.py

@@ -2,7 +2,7 @@ import collections
 import dataclasses
 import logging
 from datetime import datetime
-from typing import Dict, Iterable, List
+from typing import Any, Dict, Iterable, List, Optional
 
 from dateutil import parser
 from pydantic.fields import Field
@@ -74,15 +74,22 @@ class ClickHouseUsageConfig(ClickHouseConfig, BaseUsageConfig, EnvConfigMixin):
     options: dict = Field(default={}, description="")
     query_log_table: str = Field(default="system.query_log", exclude=True)
 
-    def get_sql_alchemy_url(self):
-        return super().get_sql_alchemy_url()
+    def get_sql_alchemy_url(
+        self,
+        uri_opts: Optional[Dict[str, Any]] = None,
+        current_db: Optional[str] = None,
+    ) -> str:
+        return super().get_sql_alchemy_url(uri_opts=uri_opts, current_db=current_db)
 
 
 @platform_name("ClickHouse")
 @config_class(ClickHouseUsageConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
+@capability(
+    SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 @dataclasses.dataclass
 class ClickHouseUsageSource(Source):
     """

datahub/ingestion/source/usage/starburst_trino_usage.py

@@ -4,7 +4,7 @@ import json
 import logging
 from datetime import datetime
 from email.utils import parseaddr
-from typing import Dict, Iterable, List, Optional
+from typing import Any, Dict, Iterable, List, Optional
 
 from dateutil import parser
 from pydantic.fields import Field
@@ -15,7 +15,9 @@ from sqlalchemy.engine import Engine
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.ingestion.api.decorators import (
+    SourceCapability,
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
@@ -58,7 +60,7 @@ AggregatedDataset = GenericAggregatedDataset[TrinoTableRef]
 
 class TrinoConnectorInfo(BaseModel):
     partitionIds: List[str]
-    truncated: Optional[bool]
+    truncated: Optional[bool] = None
 
 
 class TrinoAccessedMetadata(BaseModel):
@@ -78,7 +80,7 @@ class TrinoJoinedAccessEvent(BaseModel):
     table: Optional[str] = None
     accessed_metadata: List[TrinoAccessedMetadata]
     starttime: datetime = Field(alias="create_time")
-    endtime: Optional[datetime] = Field(alias="end_time")
+    endtime: Optional[datetime] = Field(None, alias="end_time")
 
 
 class EnvBasedSourceBaseConfig:
@@ -98,8 +100,10 @@ class TrinoUsageConfig(TrinoConfig, BaseUsageConfig, EnvBasedSourceBaseConfig):
     options: dict = Field(default={}, description="")
     database: str = Field(description="The name of the catalog from getting the usage")
 
-    def get_sql_alchemy_url(self):
-        return super().get_sql_alchemy_url()
+    def get_sql_alchemy_url(
+        self, uri_opts: Optional[Dict[str, Any]] = None, database: Optional[str] = None
+    ) -> str:
+        return super().get_sql_alchemy_url(uri_opts=uri_opts, database=database)
 
 
 @dataclasses.dataclass
@@ -110,6 +114,7 @@ class TrinoUsageReport(SourceReport):
 @platform_name("Trino")
 @config_class(TrinoUsageConfig)
 @support_status(SupportStatus.CERTIFIED)
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 @dataclasses.dataclass
 class TrinoUsageSource(Source):
     """

datahub/ingestion/source/usage/usage_common.py

@@ -12,15 +12,13 @@ from typing import (
     Optional,
     Tuple,
     TypeVar,
-    Union,
 )
 
 import pydantic
-from deprecated import deprecated
 from pydantic.fields import Field
 
 import datahub.emitter.mce_builder as builder
-from datahub.configuration.common import AllowDenyPattern
+from datahub.configuration.common import AllowDenyPattern, HiddenFromDocs
 from datahub.configuration.time_window_config import (
     BaseTimeWindowConfig,
     BucketDuration,
@@ -28,19 +26,13 @@ from datahub.configuration.time_window_config import (
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.metadata.com.linkedin.pegasus2avro.dataset import DatasetUsageStatistics
 from datahub.metadata.schema_classes import (
-    CalendarIntervalClass,
     DatasetFieldUsageCountsClass,
     DatasetUsageStatisticsClass,
     DatasetUserUsageCountsClass,
     TimeWindowSizeClass,
-    UsageAggregationClass,
-    WindowDurationClass,
 )
 from datahub.utilities.sql_formatter import format_sql_query, trim_query
-from datahub.utilities.urns.dataset_urn import DatasetUrn
-from datahub.utilities.urns.urn import guess_entity_type
 
 logger = logging.getLogger(__name__)
 
@@ -202,13 +194,13 @@ class GenericAggregatedDataset(Generic[ResourceType]):
 
 
 class BaseUsageConfig(BaseTimeWindowConfig):
-    queries_character_limit: int = Field(
+    queries_character_limit: HiddenFromDocs[int] = Field(
+        # Hidden since we don't want to encourage people to break elasticsearch.
         default=DEFAULT_QUERIES_CHARACTER_LIMIT,
         description=(
             "Total character limit for all queries in a single usage aspect."
             " Queries will be truncated to length `queries_character_limit / top_n_queries`."
         ),
-        hidden_from_docs=True,  # Don't want to encourage people to break elasticsearch
     )
 
     top_n_queries: pydantic.PositiveInt = Field(
@@ -276,6 +268,7 @@ class UsageAggregator(Generic[ResourceType]):
             user,
             query,
             fields,
+            user_email_pattern=self.config.user_email_pattern,
             count=count,
         )
 
@@ -295,60 +288,3 @@ class UsageAggregator(Generic[ResourceType]):
                     user_urn_builder=user_urn_builder,
                     queries_character_limit=self.config.queries_character_limit,
                 )
-
-
-@deprecated
-def convert_usage_aggregation_class(
-    obj: UsageAggregationClass,
-) -> MetadataChangeProposalWrapper:
-    # Legacy usage aggregation only supported dataset usage stats
-    if guess_entity_type(obj.resource) == DatasetUrn.ENTITY_TYPE:
-        aspect = DatasetUsageStatistics(
-            timestampMillis=obj.bucket,
-            eventGranularity=TimeWindowSizeClass(
-                unit=convert_window_to_interval(obj.duration)
-            ),
-            uniqueUserCount=obj.metrics.uniqueUserCount,
-            totalSqlQueries=obj.metrics.totalSqlQueries,
-            topSqlQueries=obj.metrics.topSqlQueries,
-            userCounts=(
-                [
-                    DatasetUserUsageCountsClass(
-                        user=u.user, count=u.count, userEmail=u.userEmail
-                    )
-                    for u in obj.metrics.users
-                    if u.user is not None
-                ]
-                if obj.metrics.users
-                else None
-            ),
-            fieldCounts=(
-                [
-                    DatasetFieldUsageCountsClass(fieldPath=f.fieldName, count=f.count)
-                    for f in obj.metrics.fields
-                ]
-                if obj.metrics.fields
-                else None
-            ),
-        )
-        return MetadataChangeProposalWrapper(entityUrn=obj.resource, aspect=aspect)
-    else:
-        raise Exception(
-            f"Skipping unsupported usage aggregation - invalid entity type: {obj}"
-        )
-
-
-@deprecated
-def convert_window_to_interval(window: Union[str, WindowDurationClass]) -> str:
-    if window == WindowDurationClass.YEAR:
-        return CalendarIntervalClass.YEAR
-    elif window == WindowDurationClass.MONTH:
-        return CalendarIntervalClass.MONTH
-    elif window == WindowDurationClass.WEEK:
-        return CalendarIntervalClass.WEEK
-    elif window == WindowDurationClass.DAY:
-        return CalendarIntervalClass.DAY
-    elif window == WindowDurationClass.HOUR:
-        return CalendarIntervalClass.HOUR
-    else:
-        raise Exception(f"Unsupported window duration: {window}")