acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/api/entities/structuredproperties/structuredproperties.py

@@ -1,10 +1,10 @@
 import logging
 from enum import Enum
 from pathlib import Path
-from typing import Iterable, List, Optional, Union
+from typing import Iterable, List, Optional, Type, Union
 
 import yaml
-from pydantic import StrictStr, validator
+from pydantic import Field, StrictStr, validator
 from ruamel.yaml import YAML
 
 from datahub.configuration.common import ConfigModel
@@ -48,7 +48,7 @@ VALID_ENTITY_TYPE_URNS = [
 _VALID_ENTITY_TYPES_STRING = f"Valid entity type urns are {', '.join(VALID_ENTITY_TYPE_URNS)}, etc... Ensure that the entity type is valid."
 
 
-def _validate_entity_type_urn(v: str) -> str:
+def _validate_entity_type_urn(cls: Type, v: str) -> str:
     urn = Urn.make_entity_type_urn(v)
     if urn not in VALID_ENTITY_TYPE_URNS:
         raise ValueError(
@@ -68,7 +68,7 @@ class TypeQualifierAllowedTypes(ConfigModel):
 
 class StructuredProperties(ConfigModel):
     id: Optional[str] = None
-    urn: Optional[str] = None
+    urn: Optional[str] = Field(None, validate_default=True)
     qualified_name: Optional[str] = None
     type: str
     value_entity_types: Optional[List[str]] = None

datahub/api/graphql/operation.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-from gql import gql
+from gql import GraphQLRequest
 
 from datahub.api.graphql.base import BaseApi
 
@@ -79,10 +79,12 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
         if custom_properties is not None:
             variable_values["customProperties"] = custom_properties
 
-        result = self.client.execute(
-            gql(Operation.REPORT_OPERATION_MUTATION), variable_values
+        request = GraphQLRequest(
+            Operation.REPORT_OPERATION_MUTATION, variable_values=variable_values
         )
 
+        result = self.client.execute(request)
+
         return result["reportOperation"]
 
     def query_operations(
@@ -109,12 +111,12 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
         :param partition: The partition to check the operation.
         """
 
-        result = self.client.execute(
-            gql(Operation.QUERY_OPERATIONS),
+        request = GraphQLRequest(
+            Operation.QUERY_OPERATIONS,
             variable_values={
                 "urn": urn,
                 "startTimeMillis": start_time_millis,
-                "end_time_millis": end_time_millis,
+                "endTimeMillis": end_time_millis,
                 "limit": limit,
                 "filter": self.gen_filter(
                     {
@@ -125,6 +127,8 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
                 ),
             },
         )
+
+        result = self.client.execute(request)
         if "dataset" in result and "operations" in result["dataset"]:
             operations = []
             if source_type is not None:

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/config_utils.py

@@ -11,7 +11,16 @@ import click
 import yaml
 from pydantic import BaseModel, ValidationError
 
-from datahub.cli.env_utils import get_boolean_env_variable
+from datahub.configuration.env_vars import (
+    get_gms_host,
+    get_gms_port,
+    get_gms_protocol,
+    get_gms_token,
+    get_gms_url,
+    get_skip_config,
+    get_system_client_id,
+    get_system_client_secret,
+)
 from datahub.ingestion.graph.config import DatahubClientConfig
 
 logger = logging.getLogger(__name__)
@@ -36,15 +45,15 @@ class MissingConfigError(Exception):
 
 
 def get_system_auth() -> Optional[str]:
-    system_client_id = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_ID)
-    system_client_secret = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_SECRET)
+    system_client_id = get_system_client_id()
+    system_client_secret = get_system_client_secret()
     if system_client_id is not None and system_client_secret is not None:
         return f"Basic {system_client_id}:{system_client_secret}"
     return None
 
 
 def _should_skip_config() -> bool:
-    return get_boolean_env_variable(ENV_SKIP_CONFIG, False)
+    return get_skip_config()
 
 
 def persist_raw_datahub_config(config: dict) -> None:
@@ -67,11 +76,11 @@ class DatahubConfig(BaseModel):
 
 
 def _get_config_from_env() -> Tuple[Optional[str], Optional[str]]:
-    host = os.environ.get(ENV_METADATA_HOST)
-    port = os.environ.get(ENV_METADATA_PORT)
-    token = os.environ.get(ENV_METADATA_TOKEN)
-    protocol = os.environ.get(ENV_METADATA_PROTOCOL, "http")
-    url = os.environ.get(ENV_METADATA_HOST_URL)
+    host = get_gms_host()
+    port = get_gms_port()
+    token = get_gms_token()
+    protocol = get_gms_protocol()
+    url = get_gms_url()
     if port is not None:
         url = f"{protocol}://{host}:{port}"
         return url, token
@@ -108,7 +117,6 @@ def load_client_config() -> DatahubClientConfig:
         datahub_config: DatahubClientConfig = DatahubConfig.parse_obj(
             client_config_dict
         ).gms
-
         return datahub_config
     except ValidationError as e:
         click.echo(f"Error loading your {CONDENSED_DATAHUB_CONFIG_PATH}")

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

datahub/cli/delete_cli.py

@@ -1,5 +1,6 @@
 import logging
 import random
+import sys
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from datetime import datetime
@@ -17,7 +18,6 @@ from datahub.emitter.aspect import ASPECT_MAP, TIMESERIES_ASPECT_MAP
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
 from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.graph.filters import RemovedStatusFilter
-from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 from datahub.utilities.perf_timer import PerfTimer
 from datahub.utilities.urns.urn import guess_entity_type
@@ -115,7 +115,7 @@ class DeletionResult:
     help="specifies soft/hard deletion",
 )
 @click.option("-n", "--dry-run", required=False, is_flag=True)
-@telemetry.with_telemetry()
+@upgrade.check_upgrade
 def by_registry(
     registry_id: str,
     soft: bool,
@@ -170,7 +170,7 @@ def by_registry(
 @click.option(
     "-f", "--force", required=False, is_flag=True, help="force the delete if set"
 )
-@telemetry.with_telemetry()
+@upgrade.check_upgrade
 def references(urn: str, dry_run: bool, force: bool) -> None:
     """
     Delete all references to an entity (but not the entity itself).
@@ -231,8 +231,9 @@ def references(urn: str, dry_run: bool, force: bool) -> None:
     default=3000,
     type=int,
     help="Batch size when querying for entities to un-soft delete."
-    "Maximum 10000. Large batch sizes may cause timeouts.",
+    "Maximum 5000. Large batch sizes may cause timeouts.",
 )
+@upgrade.check_upgrade
 def undo_by_filter(
     urn: Optional[str], platform: Optional[str], batch_size: int
 ) -> None:
@@ -317,6 +318,19 @@ def undo_by_filter(
     is_flag=True,
     help="Recursively delete all contained entities (only for containers and dataPlatformInstances)",
 )
+@click.option(
+    "--streaming-batch",
+    required=False,
+    is_flag=True,
+    help="Use streaming batch deletion for recursive operations. Benefit of being resumable for large hierarchies where getting all URNs at once can take a long time.",
+)
+@click.option(
+    "--streaming-batch-size",
+    required=False,
+    default=12000,
+    type=int,
+    help="Batch size for streaming batch deletion for recursive operations.",
+)
 @click.option(
     "--start-time",
     required=False,
@@ -336,7 +350,7 @@ def undo_by_filter(
     default=3000,
     type=int,
     help="Batch size when querying for entities to delete."
-    "Maximum 10000. Large batch sizes may cause timeouts.",
+    "Maximum 5000. Large batch sizes may cause timeouts.",
 )
 @click.option(
     "-n",
@@ -356,7 +370,6 @@ def undo_by_filter(
     "--workers", type=int, default=1, help="Num of workers to use for deletion."
 )
 @upgrade.check_upgrade
-@telemetry.with_telemetry()
 def by_filter(
     urn: Optional[str],
     urn_file: Optional[str],
@@ -368,6 +381,8 @@ def by_filter(
     entity_type: Optional[str],
     query: Optional[str],
     recursive: bool,
+    streaming_batch: bool,
+    streaming_batch_size: int,
     start_time: Optional[datetime],
     end_time: Optional[datetime],
     batch_size: int,
@@ -386,6 +401,7 @@ def by_filter(
         env=env,
         query=query,
         recursive=recursive,
+        streaming_batch=streaming_batch,
     )
     soft_delete_filter = _validate_user_soft_delete_flags(
         soft=soft, aspect=aspect, only_soft_deleted=only_soft_deleted
@@ -417,26 +433,27 @@ def by_filter(
     # Determine which urns to delete.
     delete_by_urn = bool(urn) and not recursive
     if urn:
-        urns = [urn]
-
         if recursive:
-            # Add children urns to the list.
-            if guess_entity_type(urn) == "dataPlatformInstance":
-                urns.extend(
-                    graph.get_urns_by_filter(
-                        platform_instance=urn,
-                        status=soft_delete_filter,
-                        batch_size=batch_size,
-                    )
-                )
-            else:
-                urns.extend(
-                    graph.get_urns_by_filter(
-                        container=urn,
-                        status=soft_delete_filter,
-                        batch_size=batch_size,
-                    )
-                )
+            _delete_urns_streaming_recursive(
+                graph=graph,
+                parent_urn=urn,
+                aspect_name=aspect,
+                soft=soft,
+                dry_run=dry_run,
+                start_time=start_time,
+                end_time=end_time,
+                workers=workers,
+                soft_delete_filter=soft_delete_filter,
+                batch_size=batch_size,
+                force=force,
+                streaming_batch_size=streaming_batch_size
+                if streaming_batch
+                else sys.maxsize,
+            )
+            return
+
+        else:
+            urns = [urn]
     elif urn_file:
         with open(urn_file, "r") as r:
             urns = []
@@ -452,6 +469,7 @@ def by_filter(
                 query=query,
                 status=soft_delete_filter,
                 batch_size=batch_size,
+                skip_cache=True,
             )
         )
         if len(urns) == 0:
@@ -557,6 +575,7 @@ def _validate_user_urn_and_filters(
     env: Optional[str],
     query: Optional[str],
     recursive: bool,
+    streaming_batch: bool,
 ) -> None:
     # Check urn / filters options.
     if urn:
@@ -592,6 +611,12 @@ def _validate_user_urn_and_filters(
             f"This will only delete {urn}. Use --recursive to delete all contained entities."
         )
 
+    # Check streaming flag.
+    if streaming_batch and not recursive:
+        raise click.UsageError(
+            "The --streaming-batch flag can only be used with --recursive."
+        )
+
 
 def _validate_user_soft_delete_flags(
     soft: bool, aspect: Optional[str], only_soft_deleted: bool
@@ -654,8 +679,8 @@ def _validate_user_aspect_flags(
 def _validate_batch_size(batch_size: int) -> None:
     if batch_size <= 0:
         raise click.UsageError("Batch size must be a positive integer.")
-    elif batch_size > 10000:
-        raise click.UsageError("Batch size cannot exceed 10,000.")
+    elif batch_size > 5000:
+        raise click.UsageError("Batch size cannot exceed 5,000.")
 
 
 def _delete_one_urn(
@@ -738,3 +763,76 @@ def _delete_one_urn(
         num_timeseries_records=ts_rows_affected,
         num_referenced_entities=referenced_entities_affected,
     )
+
+
+def _delete_urns_streaming_recursive(
+    graph: DataHubGraph,
+    parent_urn: str,
+    aspect_name: Optional[str],
+    soft: bool,
+    dry_run: bool,
+    start_time: Optional[datetime],
+    end_time: Optional[datetime],
+    workers: int,
+    soft_delete_filter: RemovedStatusFilter,
+    batch_size: int,
+    force: bool,
+    streaming_batch_size: int,
+) -> None:
+    """Streaming recursive batch deletion that processes URNs in batches."""
+
+    entity_type = guess_entity_type(parent_urn)
+    click.echo(f"Starting recursive deletion of {entity_type} {parent_urn}")
+
+    if not force and not dry_run:
+        click.confirm(
+            f"This will recursively delete {parent_urn} and all its contained entities. Do you want to continue?",
+            abort=True,
+        )
+
+    urns = []
+
+    if entity_type == "dataPlatformInstance":
+        child_urns_iter = graph.get_urns_by_filter(
+            platform_instance=parent_urn,
+            status=soft_delete_filter,
+            batch_size=batch_size,
+            # Important to skip cache so we can resume from where we left off.
+            skip_cache=True,
+        )
+    else:
+        child_urns_iter = graph.get_urns_by_filter(
+            container=parent_urn,
+            status=soft_delete_filter,
+            batch_size=batch_size,
+            # Important to skip cache so we can resume from where we left off.
+            skip_cache=True,
+        )
+
+    for child_urn in child_urns_iter:
+        urns.append(child_urn)
+        if len(urns) >= streaming_batch_size:
+            _delete_urns_parallel(
+                graph=graph,
+                urns=urns,
+                aspect_name=aspect_name,
+                soft=soft,
+                dry_run=dry_run,
+                delete_by_urn=False,
+                start_time=start_time,
+                end_time=end_time,
+                workers=workers,
+            )
+            urns = []
+    urns.append(parent_urn)
+    _delete_urns_parallel(
+        graph=graph,
+        urns=urns,
+        aspect_name=aspect_name,
+        soft=soft,
+        dry_run=dry_run,
+        delete_by_urn=False,
+        start_time=start_time,
+        end_time=end_time,
+        workers=workers,
+    )