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

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 = []
@@ -557,6 +574,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 +610,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 +678,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 +762,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,
+    )

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