From 15ed3aecf0130f946f5a7e5bdb9822caf2b305c7 Mon Sep 17 00:00:00 2001
From: Kevin Hu <6051736+kevinhu@users.noreply.github.com>
Date: Wed, 8 Dec 2021 16:09:04 -0800
Subject: [PATCH] refactor(ingest): cli deletion function (#3694)

---
 .../src/datahub/cli/delete_cli.py             | 33 ++++++++++++++++---
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/metadata-ingestion/src/datahub/cli/delete_cli.py b/metadata-ingestion/src/datahub/cli/delete_cli.py
index d212e3aeb8..a2e2527937 100644
--- a/metadata-ingestion/src/datahub/cli/delete_cli.py
+++ b/metadata-ingestion/src/datahub/cli/delete_cli.py
@@ -54,6 +54,7 @@ class DeletionResult:
             self.sample_records.extend(another_result.sample_records)
 
 
+@telemetry.with_telemetry
 def delete_for_registry(
     registry_id: str,
     soft: bool,
@@ -123,7 +124,7 @@ def delete(
         session, host = cli_utils.get_session_and_host()
         entity_type = guess_entity_type(urn=urn)
         logger.info(f"DataHub configured with {host}")
-        deletion_result: DeletionResult = delete_one_urn(
+        deletion_result: DeletionResult = delete_one_urn_cmd(
             urn,
             soft=soft,
             dry_run=dry_run,
@@ -209,7 +210,7 @@ def delete_with_filters(
         )
 
     for urn in progressbar.progressbar(urns, redirect_stdout=True):
-        one_result = delete_one_urn(
+        one_result = _delete_one_urn(
             urn,
             soft=soft,
             dry_run=dry_run,
@@ -222,8 +223,7 @@ def delete_with_filters(
     return batch_deletion_result
 
 
-@telemetry.with_telemetry
-def delete_one_urn(
+def _delete_one_urn(
     urn: str,
     soft: bool = False,
     dry_run: bool = False,
@@ -271,3 +271,28 @@ def delete_one_urn(
 
     deletion_result.end()
     return deletion_result
+
+
+@telemetry.with_telemetry
+def delete_one_urn_cmd(
+    urn: str,
+    soft: bool = False,
+    dry_run: bool = False,
+    entity_type: str = "dataset",
+    cached_session_host: Optional[Tuple[sessions.Session, str]] = None,
+    cached_emitter: Optional[rest_emitter.DatahubRestEmitter] = None,
+) -> DeletionResult:
+    """
+    Wrapper around delete_one_urn because it is also called in a loop via delete_with_filters.
+
+    This is a separate function that is called only when a single URN is deleted via the CLI.
+    """
+
+    return _delete_one_urn(
+        urn,
+        soft,
+        dry_run,
+        entity_type,
+        cached_session_host,
+        cached_emitter,
+    )