acryl-datahub 1.3.0.1rc9__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

datahub/sdk/__init__.py

@@ -28,6 +28,7 @@ from datahub.sdk.main_client import DataHubClient
 from datahub.sdk.mlmodel import MLModel
 from datahub.sdk.mlmodelgroup import MLModelGroup
 from datahub.sdk.search_filters import Filter, FilterDsl
+from datahub.sdk.tag import Tag
 
 # We want to print out the warning if people do `from datahub.sdk import X`.
 # But we don't want to print out warnings if they're doing a more direct

datahub/sdk/_all_entities.py

@@ -9,6 +9,7 @@ from datahub.sdk.dataset import Dataset
 from datahub.sdk.entity import Entity
 from datahub.sdk.mlmodel import MLModel
 from datahub.sdk.mlmodelgroup import MLModelGroup
+from datahub.sdk.tag import Tag
 
 # Base entity classes that don't have circular dependencies
 # Those that do are imported in the EntityClient where needed
@@ -22,6 +23,7 @@ ENTITY_CLASSES_LIST: List[Type[Entity]] = [
     DataJob,
     Dashboard,
     Chart,
+    Tag,
 ]
 
 # Create the mapping of entity types to classes

datahub/sdk/search_filters.py

@@ -16,6 +16,7 @@ from typing import (
 )
 
 import pydantic
+from pydantic import field_validator
 
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.pydantic_migration_helpers import (
@@ -102,7 +103,8 @@ class _EntitySubtypeFilter(_BaseFilter):
         description="The entity subtype to filter on. Can be 'Table', 'View', 'Source', etc. depending on the native platform's concepts.",
     )
 
-    @pydantic.validator("entity_subtype", pre=True)
+    @field_validator("entity_subtype", mode="before")
+    @classmethod
     def validate_entity_subtype(cls, v: str) -> List[str]:
         return [v] if not isinstance(v, list) else v
 
@@ -141,10 +143,13 @@ class _PlatformFilter(_BaseFilter):
     platform: List[str]
     # TODO: Add validator to convert string -> list of strings
 
-    @pydantic.validator("platform", each_item=True)
-    def validate_platform(cls, v: str) -> str:
+    @field_validator("platform", mode="before")
+    @classmethod
+    def validate_platform(cls, v):
         # Subtle - we use the constructor instead of the from_string method
         # because coercion is acceptable here.
+        if isinstance(v, list):
+            return [str(DataPlatformUrn(item)) for item in v]
         return str(DataPlatformUrn(v))
 
     def _build_rule(self) -> SearchFilterRule:
@@ -161,8 +166,11 @@ class _PlatformFilter(_BaseFilter):
 class _DomainFilter(_BaseFilter):
     domain: List[str]
 
-    @pydantic.validator("domain", each_item=True)
-    def validate_domain(cls, v: str) -> str:
+    @field_validator("domain", mode="before")
+    @classmethod
+    def validate_domain(cls, v):
+        if isinstance(v, list):
+            return [str(DomainUrn.from_string(item)) for item in v]
         return str(DomainUrn.from_string(v))
 
     def _build_rule(self) -> SearchFilterRule:
@@ -183,8 +191,11 @@ class _ContainerFilter(_BaseFilter):
         description="If true, only entities that are direct descendants of the container will be returned.",
     )
 
-    @pydantic.validator("container", each_item=True)
-    def validate_container(cls, v: str) -> str:
+    @field_validator("container", mode="before")
+    @classmethod
+    def validate_container(cls, v):
+        if isinstance(v, list):
+            return [str(ContainerUrn.from_string(item)) for item in v]
         return str(ContainerUrn.from_string(v))
 
     @classmethod
@@ -249,17 +260,25 @@ class _OwnerFilter(_BaseFilter):
         description="The owner to filter on. Should be user or group URNs.",
     )
 
-    @pydantic.validator("owner", each_item=True)
-    def validate_owner(cls, v: str) -> str:
-        if not v.startswith("urn:li:"):
-            raise ValueError(f"Owner must be a valid User or Group URN, got: {v}")
-        _type = guess_entity_type(v)
-        if _type == CorpUserUrn.ENTITY_TYPE:
-            return str(CorpUserUrn.from_string(v))
-        elif _type == CorpGroupUrn.ENTITY_TYPE:
-            return str(CorpGroupUrn.from_string(v))
-        else:
-            raise ValueError(f"Owner must be a valid User or Group URN, got: {v}")
+    @field_validator("owner", mode="before")
+    @classmethod
+    def validate_owner(cls, v):
+        validated = []
+        for owner in v:
+            if not owner.startswith("urn:li:"):
+                raise ValueError(
+                    f"Owner must be a valid User or Group URN, got: {owner}"
+                )
+            _type = guess_entity_type(owner)
+            if _type == CorpUserUrn.ENTITY_TYPE:
+                validated.append(str(CorpUserUrn.from_string(owner)))
+            elif _type == CorpGroupUrn.ENTITY_TYPE:
+                validated.append(str(CorpGroupUrn.from_string(owner)))
+            else:
+                raise ValueError(
+                    f"Owner must be a valid User or Group URN, got: {owner}"
+                )
+        return validated
 
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
@@ -279,17 +298,21 @@ class _GlossaryTermFilter(_BaseFilter):
         description="The glossary term to filter on. Should be glossary term URNs.",
     )
 
-    @pydantic.validator("glossary_term", each_item=True)
-    def validate_glossary_term(cls, v: str) -> str:
-        if not v.startswith("urn:li:"):
-            raise ValueError(f"Glossary term must be a valid URN, got: {v}")
-        # Validate that it's a glossary term URN
-        _type = guess_entity_type(v)
-        if _type != "glossaryTerm":
-            raise ValueError(
-                f"Glossary term must be a valid glossary term URN, got: {v}"
-            )
-        return v
+    @field_validator("glossary_term", mode="before")
+    @classmethod
+    def validate_glossary_term(cls, v):
+        validated = []
+        for term in v:
+            if not term.startswith("urn:li:"):
+                raise ValueError(f"Glossary term must be a valid URN, got: {term}")
+            # Validate that it's a glossary term URN
+            _type = guess_entity_type(term)
+            if _type != "glossaryTerm":
+                raise ValueError(
+                    f"Glossary term must be a valid glossary term URN, got: {term}"
+                )
+            validated.append(term)
+        return validated
 
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
@@ -309,15 +332,19 @@ class _TagFilter(_BaseFilter):
         description="The tag to filter on. Should be tag URNs.",
     )
 
-    @pydantic.validator("tag", each_item=True)
-    def validate_tag(cls, v: str) -> str:
-        if not v.startswith("urn:li:"):
-            raise ValueError(f"Tag must be a valid URN, got: {v}")
-        # Validate that it's a tag URN
-        _type = guess_entity_type(v)
-        if _type != "tag":
-            raise ValueError(f"Tag must be a valid tag URN, got: {v}")
-        return v
+    @field_validator("tag", mode="before")
+    @classmethod
+    def validate_tag(cls, v):
+        validated = []
+        for tag in v:
+            if not tag.startswith("urn:li:"):
+                raise ValueError(f"Tag must be a valid URN, got: {tag}")
+            # Validate that it's a tag URN
+            _type = guess_entity_type(tag)
+            if _type != "tag":
+                raise ValueError(f"Tag must be a valid tag URN, got: {tag}")
+            validated.append(tag)
+        return validated
 
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
@@ -426,7 +453,8 @@ class _Not(_BaseFilter):
 
     not_: "Filter" = pydantic.Field(alias="not")
 
-    @pydantic.validator("not_", pre=False)
+    @field_validator("not_", mode="after")
+    @classmethod
     def validate_not(cls, v: "Filter") -> "Filter":
         inner_filter = v.compile()
         if len(inner_filter) != 1:
@@ -571,7 +599,7 @@ def load_filters(obj: Any) -> Filter:
     if PYDANTIC_VERSION_2:
         return pydantic.TypeAdapter(Filter).validate_python(obj)  # type: ignore
     else:
-        return pydantic.parse_obj_as(Filter, obj)  # type: ignore
+        return pydantic.TypeAdapter(Filter).validate_python(obj)  # type: ignore
 
 
 # We need FilterDsl for two reasons:

datahub/sdk/tag.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from typing import Optional, Type
+
+from typing_extensions import Self
+
+import datahub.metadata.schema_classes as models
+from datahub.metadata.urns import TagUrn, Urn
+from datahub.sdk._shared import (
+    HasOwnership,
+    OwnersInputType,
+)
+from datahub.sdk.entity import Entity, ExtraAspectsType
+
+
+class Tag(
+    HasOwnership,
+    Entity,
+):
+    __slots__ = ()
+
+    @classmethod
+    def get_urn_type(cls) -> Type[TagUrn]:
+        return TagUrn
+
+    def __init__(
+        self,
+        *,
+        # Identity.
+        name: str,
+        # Tag properties.
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        color: Optional[str] = None,
+        # Standard aspects.
+        owners: Optional[OwnersInputType] = None,
+        extra_aspects: ExtraAspectsType = None,
+    ):
+        """Initialize a new Tag instance."""
+        urn = TagUrn(name=name)
+        super().__init__(urn)
+        self._set_extra_aspects(extra_aspects)
+
+        self._ensure_tag_props(
+            display_name=display_name or name,
+            description=description,
+            color=color,
+        )
+
+        if owners is not None:
+            self.set_owners(owners)
+
+    @classmethod
+    def _new_from_graph(cls, urn: Urn, current_aspects: models.AspectBag) -> Self:
+        assert isinstance(urn, TagUrn)
+        entity = cls(name=urn.name)
+        return entity._init_from_graph(current_aspects)
+
+    @property
+    def urn(self) -> TagUrn:
+        assert isinstance(self._urn, TagUrn)
+        return self._urn
+
+    def _ensure_tag_props(
+        self,
+        *,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+        color: Optional[str] = None,
+    ) -> models.TagPropertiesClass:
+        existing_props = self._get_aspect(models.TagPropertiesClass)
+        if existing_props is not None:
+            if display_name is not None:
+                existing_props.name = display_name
+            if description is not None:
+                existing_props.description = description
+            if color is not None:
+                existing_props.colorHex = color
+            return existing_props
+
+        return self._setdefault_aspect(
+            models.TagPropertiesClass(
+                name=display_name or self.urn.name,
+                description=description,
+                colorHex=color,
+            )
+        )
+
+    @property
+    def name(self) -> str:
+        return self.urn.name
+
+    @property
+    def display_name(self) -> str:
+        return self._ensure_tag_props().name
+
+    def set_display_name(self, display_name: str) -> None:
+        self._ensure_tag_props(display_name=display_name)
+
+    @property
+    def description(self) -> Optional[str]:
+        return self._ensure_tag_props().description
+
+    def set_description(self, description: str) -> None:
+        self._ensure_tag_props(description=description)
+
+    @property
+    def color(self) -> Optional[str]:
+        return self._ensure_tag_props().colorHex
+
+    def set_color(self, color: str) -> None:
+        self._ensure_tag_props(color=color)

datahub/secret/datahub_secret_store.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, field_validator
 
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.graph.config import DatahubClientConfig
@@ -18,8 +18,11 @@ class DataHubSecretStoreConfig(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    @validator("graph_client")
-    def check_graph_connection(cls, v: DataHubGraph) -> DataHubGraph:
+    @field_validator("graph_client", mode="after")
+    @classmethod
+    def check_graph_connection(
+        cls, v: Optional[DataHubGraph]
+    ) -> Optional[DataHubGraph]:
         if v is not None:
             v.test_connection()
         return v
@@ -63,7 +66,7 @@ class DataHubSecretStore(SecretStore):
 
     @classmethod
     def create(cls, config: Any) -> "DataHubSecretStore":
-        config = DataHubSecretStoreConfig.parse_obj(config)
+        config = DataHubSecretStoreConfig.model_validate(config)
         return cls(config)
 
     def close(self) -> None:

datahub/secret/file_secret_store.py

@@ -45,5 +45,5 @@ class FileSecretStore(SecretStore):
 
     @classmethod
     def create(cls, config: Any) -> "FileSecretStore":
-        config = FileSecretStoreConfig.parse_obj(config)
+        config = FileSecretStoreConfig.model_validate(config)
         return cls(config)

datahub/sql_parsing/schema_resolver.py

@@ -1,5 +1,6 @@
 import contextlib
 import pathlib
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Protocol, Set, Tuple
 
 from typing_extensions import TypedDict
@@ -22,6 +23,14 @@ from datahub.utilities.urns.field_paths import get_simple_field_path_from_v2_fie
 SchemaInfo = Dict[str, str]
 
 
+@dataclass
+class SchemaResolverReport:
+    """Report class for tracking SchemaResolver cache performance."""
+
+    num_schema_cache_hits: int = 0
+    num_schema_cache_misses: int = 0
+
+
 class GraphQLSchemaField(TypedDict):
     fieldPath: str
     nativeDataType: str
@@ -53,6 +62,7 @@ class SchemaResolver(Closeable, SchemaResolverInterface):
         env: str = DEFAULT_ENV,
         graph: Optional[DataHubGraph] = None,
         _cache_filename: Optional[pathlib.Path] = None,
+        report: Optional[SchemaResolverReport] = None,
     ):
         # Also supports platform with an urn prefix.
         self._platform = DataPlatformUrn(platform).platform_name
@@ -60,6 +70,7 @@ class SchemaResolver(Closeable, SchemaResolverInterface):
         self.env = env
 
         self.graph = graph
+        self.report = report
 
         # Init cache, potentially restoring from a previous run.
         shared_conn = None
@@ -132,12 +143,14 @@ class SchemaResolver(Closeable, SchemaResolverInterface):
 
         schema_info = self._resolve_schema_info(urn)
         if schema_info:
+            self._track_cache_hit()
             return urn, schema_info
 
         urn_lower = self.get_urn_for_table(table, lower=True)
         if urn_lower != urn:
             schema_info = self._resolve_schema_info(urn_lower)
             if schema_info:
+                self._track_cache_hit()
                 return urn_lower, schema_info
 
         # Our treatment of platform instances when lowercasing urns
@@ -152,8 +165,12 @@ class SchemaResolver(Closeable, SchemaResolverInterface):
         if urn_mixed not in {urn, urn_lower}:
             schema_info = self._resolve_schema_info(urn_mixed)
             if schema_info:
+                self._track_cache_hit()
                 return urn_mixed, schema_info
 
+        # Track cache miss for the final attempt
+        self._track_cache_miss()
+
         if self._prefers_urn_lower():
             return urn_lower, None
         else:
@@ -165,6 +182,16 @@ class SchemaResolver(Closeable, SchemaResolverInterface):
     def has_urn(self, urn: str) -> bool:
         return self._schema_cache.get(urn) is not None
 
+    def _track_cache_hit(self) -> None:
+        """Track a cache hit if reporting is enabled."""
+        if self.report is not None:
+            self.report.num_schema_cache_hits += 1
+
+    def _track_cache_miss(self) -> None:
+        """Track a cache miss if reporting is enabled."""
+        if self.report is not None:
+            self.report.num_schema_cache_misses += 1
+
     def _resolve_schema_info(self, urn: str) -> Optional[SchemaInfo]:
         if urn in self._schema_cache:
             return self._schema_cache[urn]
@@ -261,6 +288,8 @@ class _SchemaResolverWithExtras(SchemaResolverInterface):
             table, lower=self._base_resolver._prefers_urn_lower()
         )
         if urn in self._extra_schemas:
+            # Track cache hit for extra schemas
+            self._base_resolver._track_cache_hit()
             return urn, self._extra_schemas[urn]
         return self._base_resolver.resolve_table(table)
 

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -168,6 +168,12 @@ class QueryMetadata:
             query_subject_urns.add(upstream)
             if include_fields:
                 for column in sorted(self.column_usage.get(upstream, [])):
+                    # Skip empty column names to avoid creating invalid URNs
+                    if not column or not column.strip():
+                        logger.warning(
+                            f"Skipping empty upstream column name for query {self.query_id} on upstream {upstream}"
+                        )
+                        continue
                     query_subject_urns.add(
                         builder.make_schema_field_urn(upstream, column)
                     )
@@ -175,6 +181,15 @@ class QueryMetadata:
             query_subject_urns.add(downstream_urn)
             if include_fields:
                 for column_lineage in self.column_lineage:
+                    # Skip empty downstream columns to avoid creating invalid URNs
+                    if (
+                        not column_lineage.downstream.column
+                        or not column_lineage.downstream.column.strip()
+                    ):
+                        logger.warning(
+                            f"Skipping empty downstream column name for query {self.query_id} on downstream {downstream_urn}"
+                        )
+                        continue
                     query_subject_urns.add(
                         builder.make_schema_field_urn(
                             downstream_urn, column_lineage.downstream.column

datahub/sql_parsing/sqlglot_lineage.py

@@ -28,6 +28,7 @@ import sqlglot.optimizer.optimizer
 import sqlglot.optimizer.qualify
 import sqlglot.optimizer.qualify_columns
 import sqlglot.optimizer.unnest_subqueries
+from pydantic import field_validator
 
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.ingestion.graph.client import DataHubGraph
@@ -141,7 +142,8 @@ class DownstreamColumnRef(_ParserBaseModel):
     column_type: Optional[SchemaFieldDataTypeClass] = None
     native_column_type: Optional[str] = None
 
-    @pydantic.validator("column_type", pre=True)
+    @field_validator("column_type", mode="before")
+    @classmethod
     def _load_column_type(
         cls, v: Optional[Union[dict, SchemaFieldDataTypeClass]]
     ) -> Optional[SchemaFieldDataTypeClass]:
@@ -215,7 +217,8 @@ class SqlParsingDebugInfo(_ParserBaseModel):
     def error(self) -> Optional[Exception]:
         return self.table_error or self.column_error
 
-    @pydantic.validator("table_error", "column_error")
+    @field_validator("table_error", "column_error", mode="before")
+    @classmethod
     def remove_variables_from_error(cls, v: Optional[Exception]) -> Optional[Exception]:
         if v and v.__traceback__:
             # Remove local variables from the traceback to avoid memory leaks.

datahub/testing/check_sql_parser_result.py

@@ -60,8 +60,8 @@ def assert_sql_result_with_resolver(
     expected = SqlParsingResult.parse_raw(expected_file.read_text())
 
     full_diff = deepdiff.DeepDiff(
-        expected.dict(),
-        res.dict(),
+        expected.model_dump(),
+        res.model_dump(),
         exclude_regex_paths=[
             r"root.column_lineage\[\d+\].logic",
         ],

datahub/utilities/ingest_utils.py

@@ -48,7 +48,7 @@ def deploy_source_vars(
 
     deploy_options_raw = pipeline_config.pop("deployment", None)
     if deploy_options_raw is not None:
-        deploy_options = DeployOptions.parse_obj(deploy_options_raw)
+        deploy_options = DeployOptions.model_validate(deploy_options_raw)
 
         if name:
             logger.info(f"Overriding deployment name {deploy_options.name} with {name}")