acryl-datahub 1.1.0rc3__py3-none-any.whl → 1.1.0.1__py3-none-any.whl

datahub/api/entities/external/external_entities.py

@@ -0,0 +1,239 @@
+import logging
+from abc import abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Iterable, List, Optional, Union
+
+import cachetools
+from pydantic import BaseModel
+
+from datahub.api.entities.platformresource.platform_resource import (
+    PlatformResource,
+    PlatformResourceKey,
+)
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.urns import PlatformResourceUrn, Urn
+from datahub.utilities.search_utils import ElasticDocumentQuery
+
+logger = logging.getLogger(__name__)
+
+
+class PlatformResourceRepository:
+    def __init__(self, graph: DataHubGraph):
+        self.graph = graph
+        self.cache: cachetools.TTLCache = cachetools.TTLCache(maxsize=1000, ttl=60 * 5)
+
+    def search_by_filter(
+        self, filter: ElasticDocumentQuery, add_to_cache: bool = True
+    ) -> Iterable[PlatformResource]:
+        results = PlatformResource.search_by_filters(self.graph, filter)
+        for platform_resource in results:
+            if add_to_cache:
+                self.cache[platform_resource.id] = platform_resource
+            yield platform_resource
+
+    def create(self, platform_resource: PlatformResource) -> None:
+        platform_resource.to_datahub(self.graph)
+        self.cache[platform_resource.id] = platform_resource
+
+    def get(self, key: PlatformResourceKey) -> Optional[PlatformResource]:
+        return self.cache.get(key.id)
+
+    def delete(self, key: PlatformResourceKey) -> None:
+        self.graph.delete_entity(urn=PlatformResourceUrn(key.id).urn(), hard=True)
+        del self.cache[key.id]
+
+
+class ExternalEntityId:
+    """
+    ExternalEntityId is a unique
+    identifier for an ExternalEntity.
+    """
+
+    @abstractmethod
+    def to_platform_resource_key(self) -> PlatformResourceKey:
+        """
+        Converts the ExternalEntityId to a PlatformResourceKey.
+        """
+        pass
+
+
+class CaseSensitivity(Enum):
+    UPPER = "upper"
+    LOWER = "lower"
+    MIXED = "mixed"
+
+    @staticmethod
+    def detect_case_sensitivity(value: str) -> "CaseSensitivity":
+        if value.isupper():
+            return CaseSensitivity.UPPER
+        elif value.islower():
+            return CaseSensitivity.LOWER
+        return CaseSensitivity.MIXED
+
+    @staticmethod
+    def detect_for_many(values: List[str]) -> "CaseSensitivity":
+        """
+        Detects the case sensitivity for a list of strings.
+        Returns CaseSensitivity.MIXED if the case sensitivity is mixed.
+        """
+        if len(values) == 0:
+            return CaseSensitivity.MIXED
+
+        if all(
+            CaseSensitivity.detect_case_sensitivity(value) == CaseSensitivity.UPPER
+            for value in values
+        ):
+            return CaseSensitivity.UPPER
+        elif all(
+            CaseSensitivity.detect_case_sensitivity(value) == CaseSensitivity.LOWER
+            for value in values
+        ):
+            return CaseSensitivity.LOWER
+        return CaseSensitivity.MIXED
+
+
+class LinkedResourceSet(BaseModel):
+    """
+    A LinkedResourceSet is a set of DataHub URNs that are linked to an ExternalEntity.
+    """
+
+    urns: List[str]
+
+    def _has_conflict(self, urn: Urn) -> bool:
+        """
+        Detects if the urn is safe to add into the set
+        This is used to detect conflicts between DataHub URNs that are linked to
+        the same ExternalEntity.
+        e.g. Case sensitivity of URNs
+        Mixing tags and terms in the same set etc.
+        Return True if the urn is not safe to add into the set, else False.
+        If the urn is already in the set, we don't need to add it again, but
+        that is not a conflict.
+        """
+        if urn.urn() in self.urns:
+            return False
+
+        # Detect the entity_type of the urns in the existing set
+        detected_entity_type = None
+        for existing_urn in self.urns:
+            try:
+                parsed_urn = Urn.from_string(existing_urn)
+                entity_type = parsed_urn.entity_type
+                if detected_entity_type is None:
+                    detected_entity_type = entity_type
+                elif detected_entity_type != entity_type:
+                    logger.warning(
+                        f"Detected entity_type {detected_entity_type} is not equals to {entity_type}"
+                    )
+                    return True
+            except ValueError:
+                # Not a valid URN
+                logger.warning(f"Invalid URN {existing_urn} in LinkedResourceSet")
+                return True
+        try:
+            parsed_urn = urn
+            if (
+                detected_entity_type is not None
+                and parsed_urn.entity_type != detected_entity_type
+            ):
+                logger.warning(
+                    f"Detected entity_type {detected_entity_type} is not equals to parsed_urn's entity_type: {parsed_urn.entity_type}"
+                )
+                return True
+        except ValueError:
+            # Not a valid URN
+            logger.warning(f"Invalid URN: {urn} in LinkedResourceSet")
+            return True
+        return False
+
+    def add(self, urn: Union[str, Urn]) -> bool:
+        """
+        Adds a URN to the set.
+        Returns True if the URN was added, False if it was already in the set.
+        Raises a ValueError if the URN is in conflict with the existing set.
+        """
+        # Deduplicate the URNs if we have somehow duplicate items from concurrent runs
+        self.urns = list(set(self.urns))
+        if isinstance(urn, str):
+            urn = Urn.from_string(urn)
+        if self._has_conflict(urn):
+            raise ValueError(f"Conflict detected when adding URN {urn} to the set")
+        if urn.urn() not in self.urns:
+            self.urns.append(urn.urn())
+            return True
+        return False
+
+
+class ExternalEntity:
+    """
+    An ExternalEntity is a representation of an entity that external to DataHub
+    but could be linked to one or more DataHub entities.
+    """
+
+    @abstractmethod
+    def is_managed_by_datahub(self) -> bool:
+        """
+        Returns whether the entity is managed by DataHub.
+        """
+        pass
+
+    @abstractmethod
+    def datahub_linked_resources(self) -> LinkedResourceSet:
+        """
+        Returns the URNs of the DataHub entities linked to the external entity.
+        Empty list if no linked entities.
+        """
+        pass
+
+    @abstractmethod
+    def as_platform_resource(self) -> PlatformResource:
+        """
+        Converts the ExternalEntity to a PlatformResource.
+        """
+        pass
+
+    @abstractmethod
+    def get_id(self) -> ExternalEntityId:
+        """
+        Returns the ExternalEntityId for the ExternalEntity.
+        """
+        pass
+
+
+@dataclass
+class MissingExternalEntity(ExternalEntity):
+    id: ExternalEntityId
+
+    def is_managed_by_datahub(self) -> bool:
+        return False
+
+    def datahub_linked_resources(self) -> LinkedResourceSet:
+        return LinkedResourceSet(urns=[])
+
+    def as_platform_resource(self) -> Optional[PlatformResource]:  # type: ignore[override]
+        return None
+
+    def get_id(self) -> ExternalEntityId:
+        return self.id
+
+
+class ExternalSystem:
+    @abstractmethod
+    def exists(self, external_entity_id: ExternalEntityId) -> bool:
+        """
+        Returns whether the ExternalEntityId exists in the external system.
+        """
+        pass
+
+    @abstractmethod
+    def get(
+        self,
+        external_entity_id: ExternalEntityId,
+        platform_resource_repository: PlatformResourceRepository,
+    ) -> Optional[ExternalEntity]:
+        """
+        Returns the ExternalEntity for the ExternalEntityId.
+        Uses the platform resource repository to enrich the ExternalEntity with DataHub URNs.
+        """
+        pass

datahub/api/entities/external/external_tag.py

@@ -0,0 +1,145 @@
+"""
+External Tags Module
+
+This module provides tag types that integrate with external systems like DataHub and Unity Catalog.
+It builds on top of RestrictedText to provide sanitized, truncated tag handling with original value preservation.
+
+Classes:
+    - ExternalTag: DataHub-compatible tag with key/value parsing from URNs
+
+Example Usage:
+    # DataHub Tags
+    tag = ExternalTag.from_urn("urn:li:tag:environment:production")
+    datahub_urn = tag.get_datahub_tag  # Returns TagUrn object or string
+
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Tuple, Union
+
+from pydantic import BaseModel
+
+from datahub.api.entities.external.restricted_text import RestrictedText
+from datahub.metadata.urns import TagUrn
+
+
+class ExternalTag(BaseModel):
+    """A tag type that parses DataHub Tag URNs into key-value pairs with RestrictedText properties."""
+
+    key: RestrictedText
+    value: Optional[RestrictedText] = None
+
+    def __init__(
+        self,
+        key: Optional[Union[str, RestrictedText]] = None,
+        value: Optional[Union[str, RestrictedText]] = None,
+        **data: Any,
+    ) -> None:
+        """
+        Initialize ExternalTag from either a DataHub Tag URN or explicit key/value.
+
+        Args:
+            key: Explicit key value (optional for Pydantic initialization)
+            value: Explicit value (optional)
+            **data: Additional Pydantic data
+        """
+        if key is not None:
+            # Direct initialization with key/value
+            processed_key = (
+                RestrictedText(key) if not isinstance(key, RestrictedText) else key
+            )
+            processed_value = None
+            if value is not None:
+                processed_value = (
+                    RestrictedText(value)
+                    if not isinstance(value, RestrictedText)
+                    else value
+                )
+
+            super().__init__(
+                key=processed_key,
+                value=processed_value,
+                **data,
+            )
+        else:
+            # Standard pydantic initialization
+            super().__init__(**data)
+
+    @staticmethod
+    def _parse_tag_name(tag_name: str) -> Tuple[str, Optional[str]]:
+        """
+        Parse tag name into key and optional value.
+
+        If tag_name contains ':', split on first ':' into key:value
+        Otherwise, use entire tag_name as key with no value.
+
+        Args:
+            tag_name: The tag name portion from the URN
+
+        Returns:
+            Tuple of (key, value) where value may be None
+        """
+        if ":" in tag_name:
+            parts = tag_name.split(":", 1)  # Split on first ':' only
+            return parts[0], parts[1]
+        else:
+            return tag_name, None
+
+    def to_datahub_tag_urn(self) -> TagUrn:
+        """
+        Generate a DataHub Tag URN from the key and value.
+        This method creates the URN using the original (unprocessed) values.
+
+        Returns:
+            'urn:li:tag:key:value' if value exists, otherwise 'urn:li:tag:key'
+        """
+        if self.value is not None:
+            tag_name = f"{self.key.original}:{self.value.original}"
+        else:
+            tag_name = self.key.original
+
+        return TagUrn(name=tag_name)
+
+    @classmethod
+    def from_urn(cls, tag_urn: Union[str, "TagUrn"]) -> "ExternalTag":
+        """
+        Create an ExternalTag from a DataHub Tag URN.
+
+        Args:
+            tag_urn: DataHub Tag URN string or TagUrn object
+
+        Returns:
+            ExternalTag instance
+        """
+        if isinstance(tag_urn, str):
+            tag_urn = TagUrn.from_string(tag_urn)
+        key, value = cls._parse_tag_name(tag_urn.name)
+        return cls(key=key, value=value)
+
+    @classmethod
+    def from_key_value(cls, key: str, value: Optional[str] = None) -> "ExternalTag":
+        """
+        Create an ExternalTag from explicit key and value.
+
+        Args:
+            key: Tag key
+            value: Optional tag value
+
+        Returns:
+            ExternalTag instance
+        """
+        return cls(key=key, value=value)
+
+    def __str__(self) -> str:
+        """String representation of the tag."""
+        if self.value is not None:
+            return f"{self.key}:{self.value}"
+        else:
+            return str(self.key)
+
+    def __repr__(self) -> str:
+        if self.value is not None:
+            return f"ExternalTag(key={self.key!r}, value={self.value!r})"
+        else:
+            return f"ExternalTag(key={self.key!r})"

datahub/api/entities/external/restricted_text.py

@@ -0,0 +1,247 @@
+"""The `RestrictedText` module provides a custom Pydantic type that stores the original
+value but returns a truncated and sanitized version when accessed.
+
+Features:
+- Configurable maximum length with truncation
+- Character replacement (default replaces with underscore)
+- Preserves original value internally
+- Customizable truncation suffix
+- Compatible with both Pydantic v1 and v2
+"""
+
+from __future__ import annotations
+
+from typing import Any, ClassVar, Optional, Set, Union
+
+# Check Pydantic version and import accordingly
+try:
+    from pydantic import VERSION
+
+    PYDANTIC_V2 = int(VERSION.split(".")[0]) >= 2
+except (ImportError, AttributeError):
+    # Fallback for older versions that don't have VERSION
+    PYDANTIC_V2 = False
+
+if PYDANTIC_V2:
+    from pydantic import GetCoreSchemaHandler  # type: ignore[attr-defined]
+    from pydantic_core import core_schema
+else:
+    from pydantic.validators import str_validator
+
+
+class RestrictedTextConfig:
+    """Configuration class for RestrictedText."""
+
+    def __init__(
+        self,
+        max_length: Optional[int] = None,
+        replace_chars: Optional[Set[str]] = None,
+        replacement_char: Optional[str] = None,
+        truncation_suffix: Optional[str] = None,
+    ):
+        self.max_length = max_length
+        self.replace_chars = replace_chars
+        self.replacement_char = replacement_char
+        self.truncation_suffix = truncation_suffix
+
+
+class RestrictedText(str):
+    """A string type that stores the original value but returns a truncated and sanitized version.
+
+    This type allows you to:
+    - Set a maximum length for the displayed value
+    - Replace specific characters with a replacement character
+    - Access both the original and processed values
+
+    ```python
+    from pydantic import BaseModel
+
+    class TestModel(BaseModel):
+        # Basic usage with default settings
+        name: RestrictedText
+
+        # Custom max length and character replacement using Field
+        custom_field: RestrictedText = RestrictedText.with_config(
+            max_length=10,
+            forbidden_chars={' ', '-', '.'},
+            replacement_char='_'
+        )
+
+    # Usage example
+    model = TestModel(
+        name="This is a very long string with special characters!",
+        custom_field="hello-world.test"
+    )
+
+    print(model.name)  # Truncated and sanitized version
+    print(model.name.original)  # Original value
+    print(model.custom_field)  # "hello_worl..."
+    ```
+    """
+
+    # Default configuration
+    _default_max_length: ClassVar[Optional[int]] = 50
+    _default_replace_chars: ClassVar[Set[str]] = {" ", "\t", "\n", "\r"}
+    _default_replacement_char: ClassVar[str] = "_"
+    _default_truncation_suffix: ClassVar[str] = "..."
+
+    def __new__(cls, value: str = "") -> "RestrictedText":
+        """Create a new string instance."""
+        instance = str.__new__(cls, "")  # We'll set the display value later
+        return instance
+
+    def __init__(self, value: str = ""):
+        """Initialize the RestrictedText with a value."""
+        self.original: str = value
+        self.max_length = self._default_max_length
+        self.replace_chars = self._default_replace_chars
+        self.replacement_char = self._default_replacement_char
+        self.truncation_suffix = self._default_truncation_suffix
+
+        # Process the value
+        self._processed_value = self._process_value(value)
+
+    def _configure(
+        self,
+        max_length: Optional[int] = None,
+        replace_chars: Optional[Set[str]] = None,
+        replacement_char: Optional[str] = None,
+        truncation_suffix: Optional[str] = None,
+    ) -> "RestrictedText":
+        """Configure this instance with custom settings."""
+        if max_length is not None:
+            self.max_length = max_length
+        if replace_chars is not None:
+            self.replace_chars = replace_chars
+        if replacement_char is not None:
+            self.replacement_char = replacement_char
+        if truncation_suffix is not None:
+            self.truncation_suffix = truncation_suffix
+
+        # Reprocess the value with new configuration
+        self._processed_value = self._process_value(self.original)
+        return self
+
+    def _process_value(self, value: str) -> str:
+        """Process the value by replacing characters and truncating."""
+        # Replace specified characters
+        processed = value
+        for char in self.replace_chars:
+            processed = processed.replace(char, self.replacement_char)
+
+        # Truncate if necessary
+        if self.max_length is not None and len(processed) > self.max_length:
+            if len(self.truncation_suffix) >= self.max_length:
+                # If suffix is too long, just truncate without suffix
+                processed = processed[: self.max_length]
+            else:
+                # Truncate and add suffix
+                truncate_length = self.max_length - len(self.truncation_suffix)
+                processed = processed[:truncate_length] + self.truncation_suffix
+
+        return processed
+
+    def __str__(self) -> str:
+        """Return the processed (truncated and sanitized) value."""
+        return self._processed_value
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self._processed_value!r})"
+
+    @property
+    def processed(self) -> str:
+        """Get the processed (truncated and sanitized) value."""
+        return self._processed_value
+
+    @classmethod
+    def with_config(
+        cls,
+        max_length: Optional[int] = None,
+        forbidden_chars: Optional[Set[str]] = None,
+        replacement_char: Optional[str] = None,
+        truncation_suffix: Optional[str] = None,
+    ) -> RestrictedTextConfig:
+        """Create a configuration object for use as field default.
+
+        Args:
+            max_length: Maximum length of the processed string
+            forbidden_chars: Set of characters to replace
+            replacement_char: Character to use as replacement
+            truncation_suffix: Suffix to add when truncating
+
+        Returns:
+            A configuration object that can be used as field default
+        """
+        return RestrictedTextConfig(
+            max_length=max_length,
+            replace_chars=forbidden_chars,
+            replacement_char=replacement_char,
+            truncation_suffix=truncation_suffix,
+        )
+
+    # Pydantic v2 methods
+    if PYDANTIC_V2:
+
+        @classmethod
+        def _validate(
+            cls,
+            __input_value: Union[str, "RestrictedText"],
+            _: core_schema.ValidationInfo,
+        ) -> "RestrictedText":
+            """Validate and create a RestrictedText instance."""
+            if isinstance(__input_value, RestrictedText):
+                return __input_value
+            return cls(__input_value)
+
+        @classmethod
+        def __get_pydantic_core_schema__(
+            cls, source: type[Any], handler: GetCoreSchemaHandler
+        ) -> core_schema.CoreSchema:
+            """Get the Pydantic core schema for this type."""
+            return core_schema.with_info_after_validator_function(
+                cls._validate,
+                core_schema.str_schema(),
+                field_name=cls.__name__,
+            )
+
+    # Pydantic v1 methods
+    else:
+
+        @classmethod
+        def __get_validators__(cls):
+            """Pydantic v1 validator method."""
+            yield cls.validate
+
+        @classmethod
+        def validate(cls, v, field=None):
+            """Validate and create a RestrictedText instance for Pydantic v1."""
+            if isinstance(v, RestrictedText):
+                return v
+
+            if not isinstance(v, str):
+                # Let pydantic handle the string validation
+                v = str_validator(v)
+
+            # Create instance
+            instance = cls(v)
+
+            # Check if there's a field default that contains configuration
+            if (
+                field
+                and hasattr(field, "default")
+                and isinstance(field.default, RestrictedTextConfig)
+            ):
+                config = field.default
+                instance._configure(
+                    max_length=config.max_length,
+                    replace_chars=config.replace_chars,
+                    replacement_char=config.replacement_char,
+                    truncation_suffix=config.truncation_suffix,
+                )
+
+            return instance
+
+        @classmethod
+        def __modify_schema__(cls, field_schema):
+            """Modify the JSON schema for Pydantic v1."""
+            field_schema.update(type="string", examples=["example string"])