acryl-datahub 1.3.1__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

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/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}")