acryl-datahub 1.1.0.4rc2__py3-none-any.whl → 1.1.0.5__py3-none-any.whl

{acryl_datahub-1.1.0.4rc2.dist-info → acryl_datahub-1.1.0.5.dist-info}/entry_points.txt

@@ -39,6 +39,7 @@ datahub-business-glossary = datahub.ingestion.source.metadata.business_glossary:
 datahub-debug = datahub.ingestion.source.debug.datahub_debug:DataHubDebugSource
 datahub-gc = datahub.ingestion.source.gc.datahub_gc:DataHubGcSource
 datahub-lineage-file = datahub.ingestion.source.metadata.lineage:LineageFileSource
+datahub-mock-data = datahub.ingestion.source.mock_data.datahub_mock_data:DataHubMockDataSource
 dbt = datahub.ingestion.source.dbt.dbt_core:DBTCoreSource
 dbt-cloud = datahub.ingestion.source.dbt.dbt_cloud:DBTCloudSource
 delta-lake = datahub.ingestion.source.delta_lake:DeltaLakeSource

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.1.0.4rc2"
+__version__ = "1.1.0.5"
 
 
 def is_dev_mode() -> bool:

datahub/api/entities/dataset/dataset.py

@@ -383,7 +383,7 @@ class Dataset(StrictModel):
     urn: Optional[str] = None
     description: Optional[str] = None
     name: Optional[str] = None
-    schema_metadata: Optional[SchemaSpecification] = Field(alias="schema")
+    schema_metadata: Optional[SchemaSpecification] = Field(default=None, alias="schema")
     downstreams: Optional[List[str]] = None
     properties: Optional[Dict[str, str]] = None
     subtype: Optional[str] = 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)
@@ -482,26 +479,83 @@ def server_config() -> None:
 
 @check.command()
 @click.option(
-    "--urn", required=True, help="The urn or urn pattern (supports % for wildcard)"
+    "--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: str,
+    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)
 
-    result = graph.restore_indices(
+    graph.restore_indices(
         urn_pattern=urn,
         aspect=aspect,
         start=start,
         batch_size=batch_size,
+        file=file,
     )
-    click.echo(result)
+
+
+@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")

datahub/cli/delete_cli.py

@@ -18,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
@@ -116,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,
@@ -171,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).
@@ -234,6 +233,7 @@ def references(urn: str, dry_run: bool, force: bool) -> None:
     help="Batch size when querying for entities to un-soft delete."
     "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:
@@ -370,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],

datahub/cli/docker_check.py

@@ -1,8 +1,9 @@
 import enum
 import os
+import pathlib
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, Set
 
 import docker
 import docker.errors
@@ -13,6 +14,7 @@ from datahub.configuration.common import ExceptionWithProps
 
 # Docker seems to under-report memory allocated, so we also need a bit of buffer to account for it.
 MIN_MEMORY_NEEDED = 3.8  # GB
+MIN_DISK_SPACE_NEEDED = 12  # GB
 
 DOCKER_COMPOSE_PROJECT_NAME = os.getenv("DATAHUB_COMPOSE_PROJECT_NAME", "datahub")
 DATAHUB_COMPOSE_PROJECT_FILTER = {
@@ -37,6 +39,10 @@ class DockerLowMemoryError(Exception):
     SHOW_STACK_TRACE = False
 
 
+class DockerLowDiskSpaceError(Exception):
+    SHOW_STACK_TRACE = False
+
+
 class DockerComposeVersionError(Exception):
     SHOW_STACK_TRACE = False
 
@@ -102,6 +108,24 @@ def run_quickstart_preflight_checks(client: docker.DockerClient) -> None:
             "You can increase the memory allocated to Docker in the Docker settings."
         )
 
+    result = client.containers.run(
+        "alpine:latest",
+        "sh -c \"df -B1 / | tail -1 | awk '{print $2, $4}'\"",  # total, available
+        remove=True,
+        stdout=True,
+        stderr=True,
+    )
+
+    output = result.decode("utf-8").strip()
+    total_bytes, available_bytes = map(int, output.split())
+
+    available_gb = available_bytes / (1024**3)
+    if available_gb < MIN_DISK_SPACE_NEEDED:
+        raise DockerLowDiskSpaceError(
+            f"Total Docker disk space available {available_gb:.2f}GB is below the minimum threshold {MIN_DISK_SPACE_NEEDED}GB. "
+            "You can increase the disk space allocated to Docker in the Docker settings or free up disk space`"
+        )
+
 
 class ContainerStatus(enum.Enum):
     OK = "is ok"
@@ -126,10 +150,24 @@ class DockerContainerStatus:
 @dataclass
 class QuickstartStatus:
     containers: List[DockerContainerStatus]
+    volumes: Set[str]
+    # On moving to compose profiles, this CLI will no longer support running quickstart instances from earlier versions.
+    # While the check command can work, upgrades or
+    running_unsupported_version: bool
+
+    def __init__(
+        self,
+        containers: List[DockerContainerStatus],
+        volumes: List[str],
+        running_unsupported_version: bool = False,
+    ):
+        self.containers = containers
+        self.running_unsupported_version = running_unsupported_version
+        self.volumes = set(volumes)
 
     def errors(self) -> List[str]:
         if not self.containers:
-            return ["quickstart.sh or dev.sh is not running"]
+            return ["datahub is not running"]
 
         return [
             f"{container.name} {container.status.value}"
@@ -176,6 +214,26 @@ class QuickstartStatus:
             },
         )
 
+    def get_containers(self) -> Set[str]:
+        if self.containers:
+            return {container.name for container in self.containers}
+        else:
+            return set()
+
+
+def detect_legacy_quickstart_compose(containers: Set[str]) -> bool:
+    return "zookeeper" in containers
+
+
+def _get_services_from_compose(compose_file: str) -> Set[str]:
+    with open(compose_file) as config_file:
+        return yaml.safe_load(config_file).get("services", {}).keys()
+
+
+def _get_volumes_from_compose(compose_file: str) -> Set[str]:
+    with open(compose_file) as config_file:
+        return yaml.safe_load(config_file).get("volumes", {}).keys()
+
 
 def check_docker_quickstart() -> QuickstartStatus:
     container_statuses: List[DockerContainerStatus] = []
@@ -188,7 +246,7 @@ def check_docker_quickstart() -> QuickstartStatus:
             ignore_removed=True,
         )
         if len(containers) == 0:
-            return QuickstartStatus([])
+            return QuickstartStatus([], [], running_unsupported_version=False)
 
         # load the expected containers from the docker-compose file
         config_files = (
@@ -197,16 +255,17 @@ def check_docker_quickstart() -> QuickstartStatus:
             .split(",")
         )
 
-        # If using profiles, alternative check
+        # If using profiles, alternative check  ##TODO: Does this really work? Check mixpanel for usage of this.
         if config_files and "/profiles/" in config_files[0]:
             return check_docker_quickstart_profiles(client)
 
         all_containers = set()
         for config_file in config_files:
-            with open(config_file) as config_file:
-                all_containers.update(
-                    yaml.safe_load(config_file).get("services", {}).keys()
-                )
+            all_containers.update(_get_services_from_compose(config_file))
+
+        all_volumes = set()
+        for config_file in config_files:
+            all_volumes.update(_get_volumes_from_compose(config_file))
 
         existing_containers = set()
         # Check that the containers are running and healthy.
@@ -240,8 +299,12 @@ def check_docker_quickstart() -> QuickstartStatus:
             container_statuses.append(
                 DockerContainerStatus(missing, ContainerStatus.MISSING)
             )
-
-    return QuickstartStatus(container_statuses)
+        running_unsupported_version = detect_legacy_quickstart_compose(all_containers)
+    return QuickstartStatus(
+        containers=container_statuses,
+        volumes=list(all_volumes),
+        running_unsupported_version=running_unsupported_version,
+    )
 
 
 def check_docker_quickstart_profiles(client: docker.DockerClient) -> QuickstartStatus:
@@ -254,7 +317,7 @@ def check_docker_quickstart_profiles(client: docker.DockerClient) -> QuickstartS
         ignore_removed=True,
     )
     if len(containers) == 0:
-        return QuickstartStatus([])
+        return QuickstartStatus([], [], running_unsupported_version=False)
 
     existing_containers = set()
     # Check that the containers are running and healthy.
@@ -273,4 +336,36 @@ def check_docker_quickstart_profiles(client: docker.DockerClient) -> QuickstartS
 
         container_statuses.append(DockerContainerStatus(name, status))
 
-    return QuickstartStatus(container_statuses)
+    # TODO: Can this be handled with older verions?
+    return QuickstartStatus(
+        container_statuses, volumes=[], running_unsupported_version=False
+    )
+
+
+def check_upgrade_supported(
+    quickstart_compose_file: List[pathlib.Path], quickstart_status: QuickstartStatus
+) -> bool:
+    if (
+        quickstart_status.running_unsupported_version
+    ):  # we detected a legacy quickstart service
+        return False
+
+    if not quickstart_status.get_containers():  # no containers are running
+        return True
+
+    compose_services = set()
+    compose_volumes = set()
+
+    for compose_file in quickstart_compose_file:
+        compose_services.update(_get_services_from_compose(str(compose_file)))
+        compose_volumes.update(_get_volumes_from_compose(str(compose_file)))
+
+    # if all services and volumes are not the same, the state in the volumes may not be compatible with the new services.
+    # We are checking for containers and volumes per the compose file, not necessarily all of them being present
+    if (
+        compose_services == quickstart_status.get_containers()
+        and compose_volumes == quickstart_status.volumes
+    ):
+        return True
+    else:
+        return False