acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.2.0__py3-none-any.whl

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,
+        forbidden_chars: Optional[Set[str]] = None,
+        replacement_char: Optional[str] = None,
+        truncation_suffix: Optional[str] = None,
+    ):
+        self.max_length = max_length
+        self.forbidden_chars = forbidden_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_forbidden_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.forbidden_chars = self._default_forbidden_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,
+        forbidden_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 forbidden_chars is not None:
+            self.forbidden_chars = forbidden_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.forbidden_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,
+            forbidden_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,
+                    forbidden_chars=config.forbidden_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"])

datahub/api/entities/external/unity_catalog_external_entites.py

@@ -0,0 +1,173 @@
+# Import RestrictedText from your existing module
+# Uncomment and adjust the import path as needed:
+# from your_restricted_text_module import RestrictedText
+# The following is a list of tag constraints:
+# You can assign a maximum of 50 tags to a single securable object.
+# The maximum length of a tag key is 255 characters.
+# The maximum length of a tag value is 1000 characters.
+# The following characters are not allowed in tag keys:
+# . , - = / :
+# Tag search using the workspace search UI is supported only for tables, views, and table columns.
+# Tag search requires exact term matching.
+# https://learn.microsoft.com/en-us/azure/databricks/database-objects/tags#constraint
+from typing import Any, Dict, Optional, Set, Union
+
+from typing_extensions import ClassVar
+
+from datahub.api.entities.external.external_tag import ExternalTag
+from datahub.api.entities.external.restricted_text import RestrictedText
+
+
+class UnityCatalogTagKeyText(RestrictedText):
+    """RestrictedText configured for Unity Catalog tag keys."""
+
+    _default_max_length: ClassVar[int] = 255
+    # Unity Catalog tag keys: alphanumeric, hyphens, underscores, periods only
+    _default_forbidden_chars: ClassVar[Set[str]] = {
+        "\t",
+        "\n",
+        "\r",
+        ".",
+        ",",
+        "-",
+        "=",
+        "/",
+        ":",
+    }
+    _default_replacement_char: ClassVar[str] = "_"
+    _default_truncation_suffix: ClassVar[str] = ""  # No suffix for clean identifiers
+
+
+class UnityCatalogTagValueText(RestrictedText):
+    """RestrictedText configured for Unity Catalog tag values."""
+
+    _default_max_length: ClassVar[int] = 1000
+    # Unity Catalog tag values are more permissive but still have some restrictions
+    _default_forbidden_chars: ClassVar[Set[str]] = {"\t", "\n", "\r"}
+    _default_replacement_char: ClassVar[str] = " "
+    _default_truncation_suffix: ClassVar[str] = "..."
+
+
+class UnityCatalogTag(ExternalTag):
+    """
+    A tag type specifically designed for Unity Catalog tag restrictions.
+
+    Unity Catalog Tag Restrictions:
+    - Key: Max 127 characters, alphanumeric + hyphens, underscores, periods only
+    - Value: Max 256 characters, more permissive but no control characters
+    """
+
+    key: UnityCatalogTagKeyText
+    value: Optional[UnityCatalogTagValueText] = None
+
+    def __init__(
+        self,
+        key: Optional[Union[str, UnityCatalogTagKeyText]] = None,
+        value: Optional[Union[str, UnityCatalogTagValueText]] = None,
+        **data: Any,
+    ) -> None:
+        """
+        Initialize UnityCatalogTag 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 = (
+                UnityCatalogTagKeyText(key)
+                if not isinstance(key, UnityCatalogTagKeyText)
+                else key
+            )
+            processed_value = None
+            if value is not None:
+                processed_value = (
+                    UnityCatalogTagValueText(value)
+                    if not isinstance(value, UnityCatalogTagValueText)
+                    else value
+                )
+            # If value is an empty string, set it to None to not generater empty value in DataHub tag which results in key: tags
+            if not str(value):
+                processed_value = None
+
+            super().__init__(
+                key=processed_key,
+                value=processed_value,
+                **data,
+            )
+        else:
+            # Standard pydantic initialization
+            super().__init__(**data)
+
+    def __eq__(self, other: object) -> bool:
+        """Check equality based on key and value."""
+        if not isinstance(other, UnityCatalogTag):
+            return False
+        return str(self.key) == str(other.key) and (
+            str(self.value) if self.value else None
+        ) == (str(other.value) if other.value else None)
+
+    def __hash__(self) -> int:
+        """Make UnityCatalogTag hashable based on key and value."""
+        return hash((str(self.key), str(self.value) if self.value else None))
+
+    @classmethod
+    def from_dict(cls, tag_dict: Dict[str, Any]) -> "UnityCatalogTag":
+        """
+        Create a UnityCatalogTag from a dictionary with 'key' and optional 'value'.
+
+        Args:
+            tag_dict: Dictionary with 'key' and optional 'value' keys
+
+        Returns:
+            UnityCatalogTag instance
+        """
+        return cls(key=tag_dict["key"], value=tag_dict.get("value"))
+
+    @classmethod
+    def from_key_value(cls, key: str, value: Optional[str] = None) -> "UnityCatalogTag":
+        """
+        Create a UnityCatalogTag from explicit key and value.
+
+        Overrides the parent method to return the correct type.
+
+        Args:
+            key: Tag key
+            value: Optional tag value
+
+        Returns:
+            UnityCatalogTag instance
+        """
+        return cls(key=key, value=value)
+
+    def to_dict(self) -> Dict[str, str]:
+        """
+        Convert to dictionary format suitable for Unity Catalog API.
+
+        Returns:
+            Dictionary with 'key' and optionally 'value'
+        """
+        result: Dict[str, str] = {"key": self.key.original}
+        if self.value is not None:
+            result["value"] = self.value.original
+        return result
+
+    def to_display_dict(self) -> Dict[str, str]:
+        """
+        Convert to dictionary format showing processed values.
+
+        Returns:
+            Dictionary with processed 'key' and optional 'value'
+        """
+        result: Dict[str, str] = {"key": str(self.key)}
+        if self.value is not None:
+            result["value"] = str(self.value)
+        return result
+
+    def __repr__(self) -> str:
+        if self.value:
+            return f"UnityCatalogTag(key={self.key!r}, value={self.value!r})"
+        else:
+            return f"UnityCatalogTag(key={self.key!r})"

datahub/cli/check_cli.py

@@ -9,6 +9,7 @@ from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
 
 import click
+from tabulate import tabulate
 
 from datahub._version import __package_name__
 from datahub.cli.json_file import check_mce_file
@@ -21,7 +22,7 @@ from datahub.ingestion.run.pipeline import Pipeline
 from datahub.ingestion.sink.sink_registry import sink_registry
 from datahub.ingestion.source.source_registry import source_registry
 from datahub.ingestion.transformer.transform_registry import transform_registry
-from datahub.telemetry import telemetry
+from datahub.upgrade import upgrade
 from datahub.utilities.file_backed_collections import (
     ConnectionWrapper,
     FileBackedDict,
@@ -47,7 +48,6 @@ def check() -> None:
 @click.option(
     "--unpack-mces", default=False, is_flag=True, help="Converts MCEs into MCPs"
 )
-@telemetry.with_telemetry()
 def metadata_file(json_file: str, rewrite: bool, unpack_mces: bool) -> None:
     """Check the schema of a metadata (MCE or MCP) JSON file."""
 
@@ -105,7 +105,6 @@ def metadata_file(json_file: str, rewrite: bool, unpack_mces: bool) -> None:
     default=(),
     help="[Advanced] Paths in the deepdiff object to ignore",
 )
-@telemetry.with_telemetry()
 def metadata_diff(
     actual_file: str, expected_file: str, verbose: bool, ignore_path: List[str]
 ) -> None:
@@ -142,7 +141,6 @@ def metadata_diff(
     type=str,
     default=None,
 )
-@telemetry.with_telemetry()
 def plugins(source: Optional[str], verbose: bool) -> None:
     """List the enabled ingestion plugins."""
 
@@ -234,7 +232,7 @@ def sql_format(sql: str, platform: str) -> None:
     default=True,
     help="Run in offline mode and disable schema-aware parsing.",
 )
-@telemetry.with_telemetry()
+@upgrade.check_upgrade
 def sql_lineage(
     sql: Optional[str],
     sql_file: Optional[str],
@@ -297,7 +295,6 @@ def sql_lineage(
     type=str,
     help="the input to validate",
 )
-@telemetry.with_telemetry()
 def test_allow_deny(config: str, input: str, pattern_key: str) -> None:
     """Test input string against AllowDeny pattern in a DataHub recipe.
 
@@ -346,7 +343,6 @@ def test_allow_deny(config: str, input: str, pattern_key: str) -> None:
     type=str,
     help="The input to validate",
 )
-@telemetry.with_telemetry()
 def test_path_spec(config: str, input: str, path_spec_key: str) -> None:
     """Test input path string against PathSpec patterns in a DataHub recipe.
 
@@ -471,6 +467,7 @@ WHERE
 
 
 @check.command()
+@upgrade.check_upgrade
 def server_config() -> None:
     """Print the server config."""
     graph = get_default_graph(ClientMode.CLI)
@@ -478,3 +475,87 @@ def server_config() -> None:
     server_config = graph.get_server_config()
 
     click.echo(pprint.pformat(server_config))
+
+
+@check.command()
+@click.option(
+    "--urn", required=False, help="The urn or urn pattern (supports % for wildcard)"
+)
+@click.option("--aspect", default=None, help="Filter to a specific aspect name.")
+@click.option(
+    "--start", type=int, default=None, help="Row number of sql store to restore from."
+)
+@click.option("--batch-size", type=int, default=None, help="How many rows to restore.")
+@click.option(
+    "--file",
+    required=False,
+    type=click.Path(exists=True, dir_okay=True, readable=True),
+    help="File absolute path containing URNs (one per line) to restore indices",
+)
+@upgrade.check_upgrade
+def restore_indices(
+    urn: Optional[str],
+    aspect: Optional[str],
+    start: Optional[int],
+    batch_size: Optional[int],
+    file: Optional[str],
+) -> None:
+    """Resync metadata changes into the search and graph indices."""
+    if urn is None and file is None:
+        raise click.UsageError("Either --urn or --file must be provided")
+    graph = get_default_graph(ClientMode.CLI)
+
+    graph.restore_indices(
+        urn_pattern=urn,
+        aspect=aspect,
+        start=start,
+        batch_size=batch_size,
+        file=file,
+    )
+
+
+@check.command()
+@upgrade.check_upgrade
+def get_kafka_consumer_offsets() -> None:
+    """Get Kafka consumer offsets from the DataHub API."""
+    graph = get_default_graph(ClientMode.CLI)
+    result = graph.get_kafka_consumer_offsets()
+
+    table_data = []
+    headers = [
+        "Topic",
+        "Consumer Group",
+        "Schema",
+        "Partition",
+        "Offset",
+        "Lag",
+        "Avg Lag",
+        "Max Lag",
+        "Total Lag",
+    ]
+
+    for topic, consumers in result.items():
+        for consumer_group, schemas in consumers.items():
+            for schema, data in schemas.items():
+                metrics = data.get("metrics", {})
+                partitions = data.get("partitions", {})
+
+                for partition, partition_data in partitions.items():
+                    table_data.append(
+                        [
+                            topic,
+                            consumer_group,
+                            schema,
+                            partition,
+                            partition_data.get("offset", "N/A"),
+                            partition_data.get("lag", "N/A"),
+                            metrics.get("avgLag", "N/A"),
+                            metrics.get("maxLag", "N/A"),
+                            metrics.get("totalLag", "N/A"),
+                        ]
+                    )
+
+    if table_data:
+        click.echo(tabulate(table_data, headers=headers, tablefmt="grid"))
+    else:
+        click.echo("No Kafka consumer offset data found.")

datahub/cli/cli_utils.py

@@ -3,6 +3,7 @@ import logging
 import time
 import typing
 from datetime import datetime
+from functools import wraps
 from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import click
@@ -424,3 +425,65 @@ def ensure_has_system_metadata(
     props = metadata.properties
     props["clientId"] = datahub_version.__package_name__
     props["clientVersion"] = datahub_version.__version__
+
+
+def enable_auto_decorators(main_group: click.Group) -> None:
+    """
+    Enable automatic decorators for all click commands.
+    This wraps existing command callback functions to add upgrade and telemetry decorators.
+    """
+
+    def has_decorator(func: Any, module_pattern: str, function_pattern: str) -> bool:
+        """Check if function already has a specific decorator"""
+        if hasattr(func, "__wrapped__"):
+            current_func = func
+            while hasattr(current_func, "__wrapped__"):
+                # Check if this wrapper matches the module and function patterns
+                if (
+                    hasattr(current_func, "__module__")
+                    and module_pattern in current_func.__module__
+                    and hasattr(current_func, "__name__")
+                    and function_pattern in current_func.__name__
+                ):
+                    return True
+                current_func = current_func.__wrapped__
+        return False
+
+    def has_telemetry_decorator(func):
+        return has_decorator(func, "telemetry", "with_telemetry")
+
+    def wrap_command_callback(command_obj):
+        """Wrap a command's callback function to add decorators"""
+        if hasattr(command_obj, "callback") and command_obj.callback:
+            original_callback = command_obj.callback
+
+            # Import here to avoid circular imports
+            from datahub.telemetry import telemetry
+
+            decorated_callback = original_callback
+
+            if not has_telemetry_decorator(decorated_callback):
+                log.debug(
+                    f"Applying telemetry decorator to {original_callback.__module__}.{original_callback.__name__}"
+                )
+                decorated_callback = telemetry.with_telemetry()(decorated_callback)
+
+            # Preserve the original function's metadata
+            decorated_callback = wraps(original_callback)(decorated_callback)
+
+            command_obj.callback = decorated_callback
+
+    def wrap_group_commands(group_obj):
+        """Recursively wrap all commands in a group"""
+        if hasattr(group_obj, "commands"):
+            for _, command_obj in group_obj.commands.items():
+                if isinstance(command_obj, click.Group):
+                    # Recursively wrap sub-groups
+                    wrap_group_commands(command_obj)
+                else:
+                    # Wrap individual commands
+                    wrap_command_callback(command_obj)
+
+    wrap_group_commands(main_group)
+
+    log.debug("Auto-decorators enabled successfully")

datahub/cli/container_cli.py

@@ -3,6 +3,7 @@ import logging
 import click
 
 from datahub.ingestion.source.apply.datahub_apply import apply_association_to_container
+from datahub.upgrade import upgrade
 
 logger = logging.getLogger(__name__)
 
@@ -16,6 +17,7 @@ def container() -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--tag-urn", required=True, type=str)
+@upgrade.check_upgrade
 def tag(container_urn: str, tag_urn: str) -> None:
     """Add patch to add a tag to all datasets in a container"""
     apply_association_to_container(container_urn, tag_urn, "tag")
@@ -24,6 +26,7 @@ def tag(container_urn: str, tag_urn: str) -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--term-urn", required=True, type=str)
+@upgrade.check_upgrade
 def term(container_urn: str, term_urn: str) -> None:
     """Add patch to add a term to all datasets in a container"""
     apply_association_to_container(container_urn, term_urn, "term")
@@ -32,6 +35,7 @@ def term(container_urn: str, term_urn: str) -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--owner-urn", required=True, type=str)
+@upgrade.check_upgrade
 def owner(container_urn: str, owner_urn: str) -> None:
     """Add patch to add a owner to all datasets in a container"""
     apply_association_to_container(container_urn, owner_urn, "owner")
@@ -40,6 +44,7 @@ def owner(container_urn: str, owner_urn: str) -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--domain-urn", required=True, type=str)
+@upgrade.check_upgrade
 def domain(container_urn: str, domain_urn: str) -> None:
     """Add patch to add a domain to all datasets in a container"""
     apply_association_to_container(container_urn, domain_urn, "domain")