acryl-datahub 0.15.0.1rc12__py3-none-any.whl → 0.15.0.1rc13__py3-none-any.whl

{acryl_datahub-0.15.0.1rc12.dist-info → acryl_datahub-0.15.0.1rc13.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datahub/__init__.py,sha256=LdmQMvZSZxJJKW3u06itY2EkTfgOjWTGkJHo9YvmkV0,577
+datahub/__init__.py,sha256=kKM5imQ7UziWDSMvn1Ic5ZENvcshwalM2y2qGjZxUHY,577
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
 datahub/entrypoints.py,sha256=3-qSfXAx3Z0FEkBV5tlO8fQr4xk4ySeDRMVTpS5Xd6A,7793
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -169,7 +169,7 @@ datahub/ingestion/graph/config.py,sha256=3b_Gxa5wcBnphP63bBiAFdWS7PJhUHRE1WZL_q4
 datahub/ingestion/graph/connections.py,sha256=9462L0ZWGKURyypAln25eMPhK3pcufBar9tNDoqspXs,741
 datahub/ingestion/graph/filters.py,sha256=UeUZQHoimavIYx-jXLA0WGkOUe10TaO8uEZkfa-QgNE,6188
 datahub/ingestion/reporting/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py,sha256=rrnlgptYF3YkxWlLYpkLm3mgrmzHcy6AwTHUG18bKVA,8373
+datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py,sha256=O2SGDU2_qMtyr_1BH9-WkNOojFWig2z4O3M21nTRo70,9908
 datahub/ingestion/reporting/file_reporter.py,sha256=tiWukmMxHrTQI3rOAumsq6lRlw8T6spqpS6XBDYnrZU,1640
 datahub/ingestion/reporting/reporting_provider_registry.py,sha256=jTYSh3T4sensjnHQfPLiIcbA2dG8w0px9ghChAJjGdU,310
 datahub/ingestion/run/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -303,9 +303,9 @@ datahub/ingestion/source/fivetran/fivetran_log_api.py,sha256=EAak3hJpe75WZSgz6wP
 datahub/ingestion/source/fivetran/fivetran_query.py,sha256=vLrTj7e-0NxZ2U4bWTB57pih42WirqPlUvwtIRfStlQ,5275
 datahub/ingestion/source/gc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/gc/datahub_gc.py,sha256=WOg3yIaNmwdbSTwytKeSfIUihsM7FMYBip9u2Dnwk3c,12849
-datahub/ingestion/source/gc/dataprocess_cleanup.py,sha256=9brJW_HVrxJk1kAP20M7flmgYOMemOmaEl2zheWFW3c,17105
+datahub/ingestion/source/gc/dataprocess_cleanup.py,sha256=86Tm3NNWMf0xM4TklNIEeNOjEingKpYy-XvCPeaAb4k,17125
 datahub/ingestion/source/gc/execution_request_cleanup.py,sha256=sZbdkg3MuPVGf8eeeRg_2khGMZ01QoH4dgJiTxf7Srg,9813
-datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py,sha256=wRnRaIVUG483tY4nyDkEn6Xi2RL5MjrVvoCoZimqwSg,7514
+datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py,sha256=LvDGTaAaI-T0OZ3fkaFwipLdzPePunuSVWoEuSBsfEM,11099
 datahub/ingestion/source/gcs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/gcs/gcs_source.py,sha256=iwvj4JwjyVWRP1Vq106sUtQhh0GuOYVSu9zCa1wCZN0,6189
 datahub/ingestion/source/gcs/gcs_utils.py,sha256=_78KM863XXgkVLmZLtYGF5PJNnZas1go-XRtOq-79lo,1047
@@ -566,8 +566,8 @@ datahub/lite/lite_registry.py,sha256=bpH0kasP-LtwwUFNA2QsOIehfekAYfJtN-AkQLmSWnw
 datahub/lite/lite_server.py,sha256=p9Oa2nNs65mqcssSIVOr7VOzWqfVstz6ZQEdT4f82S0,1949
 datahub/lite/lite_util.py,sha256=pgBpT3vTO1YCQ2njZRNyicSkHYeEmQCt41BaXU8WvMo,4503
 datahub/metadata/__init__.py,sha256=AjhXPjI6cnpdcrBRrE5gOWo15vv2TTl2ctU4UAnUN7A,238
-datahub/metadata/_schema_classes.py,sha256=NbZUezNRH72XQUl4i_DlV-oRT4KzEBYFYcHcwGqXq9A,962516
-datahub/metadata/schema.avsc,sha256=Ulqzumt0EK7nD_OATi0hbCgw42ngoenja9SXWWsobIk,728543
+datahub/metadata/_schema_classes.py,sha256=IAWpWPxOeGmvmc96dapE0CySk1Rikbh-YieT-K9YTMY,964636
+datahub/metadata/schema.avsc,sha256=CeVb_Z7k0e5kmeqDUXUW7JDL6KSKBCdfAZzqRI_mLZo,729869
 datahub/metadata/schema_classes.py,sha256=X5Jl5EaSxyHdXOQv14pJ5WkQALun4MRpJ4q12wVFE18,1299
 datahub/metadata/urns.py,sha256=nfrCTExR-k2P9w272WVtWSN3xW1VUJngPwP3xnvULjU,1217
 datahub/metadata/_urns/__init__.py,sha256=cOF3GHMDgPhmbLKbN02NPpuLGHSu0qNgQyBRv08eqF0,243
@@ -581,7 +581,7 @@ datahub/metadata/com/linkedin/pegasus2avro/access/token/__init__.py,sha256=P9M7N
 datahub/metadata/com/linkedin/pegasus2avro/assertion/__init__.py,sha256=PgK5O-6pVRaEcvmwXAsSkwRLe8NjGiLH8AVBXeArqK8,5751
 datahub/metadata/com/linkedin/pegasus2avro/businessattribute/__init__.py,sha256=N8kO-eUi0_Rt7weizIExxlnJ2_kZRtPrZLWCC1xtDMA,653
 datahub/metadata/com/linkedin/pegasus2avro/chart/__init__.py,sha256=RNyyHLBNp_fxgFcBOLWO2UsXR1ofD_JczcBdPEQSusg,848
-datahub/metadata/com/linkedin/pegasus2avro/common/__init__.py,sha256=yD2OyoQhMT3KnvxRyzPXiHsyeH-wHG1NBlLn64iCE4A,5333
+datahub/metadata/com/linkedin/pegasus2avro/common/__init__.py,sha256=ukX0VnveTrMx9G6uDaTkuk4Z2kxXr2hUK8srZuRPxj0,5520
 datahub/metadata/com/linkedin/pegasus2avro/common/fieldtransformer/__init__.py,sha256=FN63vLiB3FCmIRqBjTA-0Xt7M6i7h5NhaVzbA1ysv18,396
 datahub/metadata/com/linkedin/pegasus2avro/connection/__init__.py,sha256=qRtw-dB14pzVzgQ0pDK8kyBplNdpRxVKNj4D70e_FqI,564
 datahub/metadata/com/linkedin/pegasus2avro/container/__init__.py,sha256=3yWt36KqDKFhRc9pzvt0AMnbMTlhKurGvT3BUvc25QU,510
@@ -705,7 +705,7 @@ datahub/metadata/schemas/DataHubViewInfo.avsc,sha256=U3fBIoG9ietLUpOknfQGNekqBdP
 datahub/metadata/schemas/DataHubViewKey.avsc,sha256=p53axIdSVbubo3r23Vpsed7NqRcQBMGveVikEHAVAok,424
 datahub/metadata/schemas/DataJobInfo.avsc,sha256=--obUbt_4X2paB39EeRKP13sBSiK-r0nq070EamoV1w,7212
 datahub/metadata/schemas/DataJobInputOutput.avsc,sha256=H1O8eAzZV34tvULdu67iBSWkdn08rt7wS208b8Nisbk,15268
-datahub/metadata/schemas/DataJobKey.avsc,sha256=_fSAQDgP_UPtZfqAZPhJmsHxxltuMh9btgw20z4R6Xk,1555
+datahub/metadata/schemas/DataJobKey.avsc,sha256=4F3myS-O6n7AlUqTvCkMSFvsYAjVhUq6uaQVbqLoYdM,1583
 datahub/metadata/schemas/DataPlatformInfo.avsc,sha256=WGPFumBNHbR75vsLrivnRCbBc8vSCuxDw2UlylMieh4,2686
 datahub/metadata/schemas/DataPlatformInstance.avsc,sha256=SNd3v_YyyLaDflv8Rd5cQR9GrVuky_cDTkYM6FqJiM8,1058
 datahub/metadata/schemas/DataPlatformInstanceKey.avsc,sha256=sXUV5EMT6N-x8d6s8ebcJ5JdFIOsJCtiiU5Jtm-ncIk,800
@@ -721,6 +721,7 @@ datahub/metadata/schemas/DataProcessInstanceRunEvent.avsc,sha256=zwTYULEnpMbqwkL
 datahub/metadata/schemas/DataProcessKey.avsc,sha256=mY1BDiEYo8RchI9DckQEz9Vks5Ibt2RdWZU8OYGnrHA,2240
 datahub/metadata/schemas/DataProductKey.avsc,sha256=tcdQNWk3pLA3xZzOnHvZuq2u4SQuk2YcAlsxE8CcEeU,621
 datahub/metadata/schemas/DataProductProperties.avsc,sha256=nYEK6JgpTprU0iZaqWLZsBGYJLkh6HCi1qCu-wbYhvM,6925
+datahub/metadata/schemas/DataTransformLogic.avsc,sha256=wDng1GK9znVoK0INHGiSCSa-AH5MrDkVdMzz4wOWmrY,2011
 datahub/metadata/schemas/DataTypeInfo.avsc,sha256=MCjzal71P8uIXZg161LrU8rZTJocZeizK-YxYA0Det0,704
 datahub/metadata/schemas/DataTypeKey.avsc,sha256=Gs5uc_azwg10e36ZbwDTFQMevr0IfiFvJoEGHRzEilw,546
 datahub/metadata/schemas/DatahubIngestionCheckpoint.avsc,sha256=m2Zyrx3ZWDc5gHuwbmBSRJ3JN4NFkpUhDEKM2Yeuqrw,5681
@@ -982,8 +983,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-0.15.0.1rc12.dist-info/METADATA,sha256=w8H0vrzaAbeZnX_mqNopX7I929V_AYXeWGKidlrUrE8,173444
-acryl_datahub-0.15.0.1rc12.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-acryl_datahub-0.15.0.1rc12.dist-info/entry_points.txt,sha256=xnPSPLK3bJGADxe4TDS4wL4u0FT_PGlahDa-ENYdYCQ,9512
-acryl_datahub-0.15.0.1rc12.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-0.15.0.1rc12.dist-info/RECORD,,
+acryl_datahub-0.15.0.1rc13.dist-info/METADATA,sha256=KnCOYV5Kg855hgL3B3zmYHzPnXVeMoZYf_3ScEj1cyA,173444
+acryl_datahub-0.15.0.1rc13.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+acryl_datahub-0.15.0.1rc13.dist-info/entry_points.txt,sha256=xnPSPLK3bJGADxe4TDS4wL4u0FT_PGlahDa-ENYdYCQ,9512
+acryl_datahub-0.15.0.1rc13.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-0.15.0.1rc13.dist-info/RECORD,,

datahub/__init__.py

@@ -3,7 +3,7 @@ import warnings
 
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "0.15.0.1rc12"
+__version__ = "0.15.0.1rc13"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py

@@ -146,12 +146,55 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
             aspect_value=source_info_aspect,
         )
 
+    @staticmethod
+    def _convert_sets_to_lists(obj: Any) -> Any:
+        """
+        Recursively converts all sets to lists in a Python object.
+        Works with nested dictionaries, lists, and sets.
+
+        Args:
+            obj: Any Python object that might contain sets
+
+        Returns:
+            The object with all sets converted to lists
+        """
+        if isinstance(obj, dict):
+            return {
+                key: DatahubIngestionRunSummaryProvider._convert_sets_to_lists(value)
+                for key, value in obj.items()
+            }
+        elif isinstance(obj, list):
+            return [
+                DatahubIngestionRunSummaryProvider._convert_sets_to_lists(element)
+                for element in obj
+            ]
+        elif isinstance(obj, set):
+            return [
+                DatahubIngestionRunSummaryProvider._convert_sets_to_lists(element)
+                for element in obj
+            ]
+        elif isinstance(obj, tuple):
+            return tuple(
+                DatahubIngestionRunSummaryProvider._convert_sets_to_lists(element)
+                for element in obj
+            )
+        else:
+            return obj
+
     def _get_recipe_to_report(self, ctx: PipelineContext) -> str:
         assert ctx.pipeline_config
         if not self.report_recipe or not ctx.pipeline_config.get_raw_dict():
             return ""
         else:
-            return json.dumps(redact_raw_config(ctx.pipeline_config.get_raw_dict()))
+            redacted_recipe = redact_raw_config(ctx.pipeline_config.get_raw_dict())
+            # This is required otherwise json dumps will fail
+            # with a TypeError: Object of type set is not JSON serializable
+            converted_recipe = (
+                DatahubIngestionRunSummaryProvider._convert_sets_to_lists(
+                    redacted_recipe
+                )
+            )
+            return json.dumps(converted_recipe)
 
     def _emit_aspect(self, entity_urn: Urn, aspect_value: _Aspect) -> None:
         self.sink.write_record_async(

datahub/ingestion/source/gc/dataprocess_cleanup.py

@@ -167,7 +167,7 @@ class DataJobEntity:
 class DataProcessCleanupReport(SourceReport):
     num_aspects_removed: int = 0
     num_aspect_removed_by_type: TopKDict[str, int] = field(default_factory=TopKDict)
-    sample_removed_aspects_by_type: TopKDict[str, LossyList[str]] = field(
+    sample_soft_deleted_aspects_by_type: TopKDict[str, LossyList[str]] = field(
         default_factory=TopKDict
     )
     num_data_flows_found: int = 0
@@ -286,9 +286,9 @@ class DataProcessCleanup:
         self.report.num_aspect_removed_by_type[type] = (
             self.report.num_aspect_removed_by_type.get(type, 0) + 1
         )
-        if type not in self.report.sample_removed_aspects_by_type:
-            self.report.sample_removed_aspects_by_type[type] = LossyList()
-        self.report.sample_removed_aspects_by_type[type].append(urn)
+        if type not in self.report.sample_soft_deleted_aspects_by_type:
+            self.report.sample_soft_deleted_aspects_by_type[type] = LossyList()
+        self.report.sample_soft_deleted_aspects_by_type[type].append(urn)
 
         if self.dry_run:
             logger.info(

datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py

@@ -1,9 +1,10 @@
 import logging
 import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import List, Optional
+from threading import Lock
+from typing import Dict, Iterable, List, Optional
 
 from pydantic import Field
 
@@ -18,12 +19,28 @@ from datahub.utilities.urns._urn_base import Urn
 
 logger = logging.getLogger(__name__)
 
+QUERY_QUERY_ENTITY = """
+query listQueries($input: ScrollAcrossEntitiesInput!) {
+  scrollAcrossEntities(input: $input) {
+    nextScrollId
+    count
+    searchResults {
+      entity {
+        ... on QueryEntity {
+          urn
+        }
+      }
+    }
+  }
+}
+"""
+
 
 class SoftDeletedEntitiesCleanupConfig(ConfigModel):
     enabled: bool = Field(
         default=True, description="Whether to do soft deletion cleanup."
     )
-    retention_days: Optional[int] = Field(
+    retention_days: int = Field(
         10,
         description="Number of days to retain metadata in DataHub",
     )
@@ -62,23 +79,30 @@ class SoftDeletedEntitiesCleanupConfig(ConfigModel):
         default=None,
         description="Query to filter entities",
     )
+
     limit_entities_delete: Optional[int] = Field(
         25000, description="Max number of entities to delete."
     )
 
-    runtime_limit_seconds: Optional[int] = Field(
-        None,
+    futures_max_at_time: int = Field(
+        1000, description="Max number of futures to have at a time."
+    )
+
+    runtime_limit_seconds: int = Field(
+        7200,  # 2 hours by default
         description="Runtime limit in seconds",
     )
 
 
 @dataclass
 class SoftDeletedEntitiesReport(SourceReport):
-    num_soft_deleted_entity_removed: int = 0
-    num_soft_deleted_entity_removed_by_type: TopKDict[str, int] = field(
-        default_factory=TopKDict
-    )
-    sample_soft_deleted_removed_aspects_by_type: TopKDict[str, LossyList[str]] = field(
+    num_queries_found: int = 0
+    num_soft_deleted_entity_processed: int = 0
+    num_soft_deleted_retained_due_to_age: int = 0
+    num_soft_deleted_entity_removal_started: int = 0
+    num_hard_deleted: int = 0
+    num_hard_deleted_by_type: TopKDict[str, int] = field(default_factory=TopKDict)
+    sample_hard_deleted_aspects_by_type: TopKDict[str, LossyList[str]] = field(
         default_factory=TopKDict
     )
 
@@ -103,48 +127,53 @@ class SoftDeletedEntitiesCleanup:
         self.config = config
         self.report = report
         self.dry_run = dry_run
+        self.start_time = 0.0
+        self._report_lock: Lock = Lock()
+        self.last_print_time = 0.0
+
+    def _increment_retained_count(self) -> None:
+        """Thread-safe method to update report fields"""
+        with self._report_lock:
+            self.report.num_soft_deleted_retained_due_to_age += 1
+
+    def _increment_removal_started_count(self) -> None:
+        """Thread-safe method to update report fields"""
+        with self._report_lock:
+            self.report.num_soft_deleted_entity_removal_started += 1
+
+    def _update_report(self, urn: str, entity_type: str) -> None:
+        """Thread-safe method to update report fields"""
+        with self._report_lock:
+            self.report.num_hard_deleted += 1
+
+            current_count = self.report.num_hard_deleted_by_type.get(entity_type, 0)
+            self.report.num_hard_deleted_by_type[entity_type] = current_count + 1
+            if entity_type not in self.report.sample_hard_deleted_aspects_by_type:
+                self.report.sample_hard_deleted_aspects_by_type[
+                    entity_type
+                ] = LossyList()
+            self.report.sample_hard_deleted_aspects_by_type[entity_type].append(urn)
 
     def delete_entity(self, urn: str) -> None:
         assert self.ctx.graph
 
         entity_urn = Urn.from_string(urn)
-        self.report.num_soft_deleted_entity_removed += 1
-        self.report.num_soft_deleted_entity_removed_by_type[entity_urn.entity_type] = (
-            self.report.num_soft_deleted_entity_removed_by_type.get(
-                entity_urn.entity_type, 0
-            )
-            + 1
-        )
-        if (
-            entity_urn.entity_type
-            not in self.report.sample_soft_deleted_removed_aspects_by_type
-        ):
-            self.report.sample_soft_deleted_removed_aspects_by_type[
-                entity_urn.entity_type
-            ] = LossyList()
-        self.report.sample_soft_deleted_removed_aspects_by_type[
-            entity_urn.entity_type
-        ].append(urn)
-
         if self.dry_run:
             logger.info(
                 f"Dry run is on otherwise it would have deleted {urn} with hard deletion"
             )
             return
-
+        self._increment_removal_started_count()
         self.ctx.graph.delete_entity(urn=urn, hard=True)
         self.ctx.graph.delete_references_to_urn(
             urn=urn,
             dry_run=False,
         )
+        self._update_report(urn, entity_urn.entity_type)
 
     def delete_soft_deleted_entity(self, urn: str) -> None:
         assert self.ctx.graph
 
-        if self.config.retention_days is None:
-            logger.info("Retention days is not set, skipping soft delete cleanup")
-            return
-
         retention_time = (
             int(datetime.now(timezone.utc).timestamp())
             - self.config.retention_days * 24 * 60 * 60
@@ -157,15 +186,85 @@ class SoftDeletedEntitiesCleanup:
             ]["created"]["time"] < (retention_time * 1000):
                 logger.debug(f"Hard deleting {urn}")
                 self.delete_entity(urn)
+            else:
+                self._increment_retained_count()
+
+    def _print_report(self) -> None:
+        time_taken = round(time.time() - self.last_print_time, 1)
+        # Print report every 2 minutes
+        if time_taken > 120:
+            self.last_print_time = time.time()
+            logger.info(f"\n{self.report.as_string()}")
+
+    def _process_futures(self, futures: Dict[Future, str]) -> Dict[Future, str]:
+        done, not_done = wait(futures, return_when=FIRST_COMPLETED)
+        futures = {future: urn for future, urn in futures.items() if future in not_done}
+
+        for future in done:
+            self._print_report()
+            if future.exception():
+                logger.error(
+                    f"Failed to delete entity {futures[future]}: {future.exception()}"
+                )
+                self.report.failure(
+                    f"Failed to delete entity {futures[future]}",
+                    exc=future.exception(),
+                )
+            self.report.num_soft_deleted_entity_processed += 1
+            if (
+                self.report.num_soft_deleted_entity_processed % self.config.batch_size
+                == 0
+            ):
+                if self.config.delay:
+                    logger.debug(
+                        f"Sleeping for {self.config.delay} seconds before further processing batch"
+                    )
+                    time.sleep(self.config.delay)
+        return futures
 
-    def cleanup_soft_deleted_entities(self) -> None:
-        if not self.config.enabled:
-            return
+    def _get_soft_deleted_queries(self) -> Iterable[str]:
         assert self.ctx.graph
-        start_time = time.time()
-
-        deleted_count_retention = 0
-        urns = self.ctx.graph.get_urns_by_filter(
+        scroll_id: Optional[str] = None
+        while True:
+            try:
+                result = self.ctx.graph.execute_graphql(
+                    QUERY_QUERY_ENTITY,
+                    {
+                        "input": {
+                            "types": ["QUERY"],
+                            "query": "*",
+                            "scrollId": scroll_id if scroll_id else None,
+                            "count": self.config.batch_size,
+                            "orFilters": [
+                                {
+                                    "and": [
+                                        {
+                                            "field": "removed",
+                                            "values": ["true"],
+                                            "condition": "EQUAL",
+                                        }
+                                    ]
+                                }
+                            ],
+                        }
+                    },
+                )
+            except Exception as e:
+                self.report.failure(
+                    f"While trying to get queries with {scroll_id}", exc=e
+                )
+                break
+            scroll_across_entities = result.get("scrollAcrossEntities")
+            if not scroll_across_entities:
+                break
+            scroll_id = scroll_across_entities.get("nextScrollId")
+            self.report.num_queries_found += scroll_across_entities.get("count")
+            for query in scroll_across_entities.get("searchResults"):
+                yield query["entity"]["urn"]
+
+    def _get_urns(self) -> Iterable[str]:
+        assert self.ctx.graph
+        yield from self.ctx.graph.get_urns_by_filter(
             entity_types=self.config.entity_types,
             platform=self.config.platform,
             env=self.config.env,
@@ -173,52 +272,41 @@ class SoftDeletedEntitiesCleanup:
             status=RemovedStatusFilter.ONLY_SOFT_DELETED,
             batch_size=self.config.batch_size,
         )
+        yield from self._get_soft_deleted_queries()
+
+    def cleanup_soft_deleted_entities(self) -> None:
+        if not self.config.enabled:
+            return
+        self.start_time = time.time()
 
-        futures = {}
+        futures: Dict[Future, str] = dict()
         with ThreadPoolExecutor(max_workers=self.config.max_workers) as executor:
-            num_urns_submitted = 0
-            for urn in urns:
-                num_urns_submitted += 1
+            for urn in self._get_urns():
+                self._print_report()
+                while len(futures) >= self.config.futures_max_at_time:
+                    futures = self._process_futures(futures)
                 if (
                     self.config.limit_entities_delete
-                    and num_urns_submitted > self.config.limit_entities_delete
+                    and self.report.num_hard_deleted > self.config.limit_entities_delete
                 ):
                     logger.info(
-                        f"Limit of {self.config.limit_entities_delete} entities reached. Stopping"
+                        f"Limit of {self.config.limit_entities_delete} entities reached. Stopped adding more."
                     )
                     break
                 if (
                     self.config.runtime_limit_seconds
-                    and time.time() - start_time > self.config.runtime_limit_seconds
+                    and time.time() - self.start_time
+                    > self.config.runtime_limit_seconds
                 ):
                     logger.info(
-                        f"Runtime limit of {self.config.runtime_limit_seconds} seconds reached. Stopping"
+                        f"Runtime limit of {self.config.runtime_limit_seconds} seconds reached. Not submitting more futures."
                     )
                     break
 
                 future = executor.submit(self.delete_soft_deleted_entity, urn)
                 futures[future] = urn
 
-            if not futures:
-                return
-            for future in as_completed(futures):
-                if future.exception():
-                    logger.error(
-                        f"Failed to delete entity {futures[future]}: {future.exception()}"
-                    )
-                    self.report.failure(
-                        f"Failed to delete entity {futures[future]}",
-                        exc=future.exception(),
-                    )
-                deleted_count_retention += 1
-
-                if deleted_count_retention % self.config.batch_size == 0:
-                    logger.info(
-                        f"Processed {deleted_count_retention} soft deleted entity and deleted {self.report.num_soft_deleted_entity_removed} entities so far"
-                    )
-
-                    if self.config.delay:
-                        logger.debug(
-                            f"Sleeping for {self.config.delay} seconds before getting next batch"
-                        )
-                        time.sleep(self.config.delay)
+            logger.info(f"Waiting for {len(futures)} futures to complete")
+            while len(futures) > 0:
+                self._print_report()
+                futures = self._process_futures(futures)

datahub/metadata/_schema_classes.py

@@ -4053,6 +4053,60 @@ class DataPlatformInstanceClass(_Aspect):
         self._inner_dict['instance'] = value
     
     
+class DataTransformClass(DictWrapper):
+    """Information about a transformation. It may be a query,"""
+    
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.common.DataTransform")
+    def __init__(self,
+        queryStatement: Union[None, "QueryStatementClass"]=None,
+    ):
+        super().__init__()
+        
+        self.queryStatement = queryStatement
+    
+    def _restore_defaults(self) -> None:
+        self.queryStatement = self.RECORD_SCHEMA.fields_dict["queryStatement"].default
+    
+    
+    @property
+    def queryStatement(self) -> Union[None, "QueryStatementClass"]:
+        """The data transform may be defined by a query statement"""
+        return self._inner_dict.get('queryStatement')  # type: ignore
+    
+    @queryStatement.setter
+    def queryStatement(self, value: Union[None, "QueryStatementClass"]) -> None:
+        self._inner_dict['queryStatement'] = value
+    
+    
+class DataTransformLogicClass(_Aspect):
+    """Information about a Query against one or more data assets (e.g. Tables or Views)."""
+
+
+    ASPECT_NAME = 'dataTransformLogic'
+    ASPECT_INFO = {}
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.common.DataTransformLogic")
+
+    def __init__(self,
+        transforms: List["DataTransformClass"],
+    ):
+        super().__init__()
+        
+        self.transforms = transforms
+    
+    def _restore_defaults(self) -> None:
+        self.transforms = list()
+    
+    
+    @property
+    def transforms(self) -> List["DataTransformClass"]:
+        """List of transformations applied"""
+        return self._inner_dict.get('transforms')  # type: ignore
+    
+    @transforms.setter
+    def transforms(self, value: List["DataTransformClass"]) -> None:
+        self._inner_dict['transforms'] = value
+    
+    
 class DeprecationClass(_Aspect):
     """Deprecation status of an entity"""
 
@@ -14624,7 +14678,7 @@ class DataJobKeyClass(_Aspect):
 
 
     ASPECT_NAME = 'dataJobKey'
-    ASPECT_INFO = {'keyForEntity': 'dataJob', 'entityCategory': '_unset_', 'entityAspects': ['datahubIngestionRunSummary', 'datahubIngestionCheckpoint', 'domains', 'deprecation', 'versionInfo', 'dataJobInfo', 'dataJobInputOutput', 'editableDataJobProperties', 'ownership', 'status', 'globalTags', 'browsePaths', 'glossaryTerms', 'institutionalMemory', 'dataPlatformInstance', 'browsePathsV2', 'structuredProperties', 'forms', 'subTypes', 'incidentsSummary', 'testResults']}
+    ASPECT_INFO = {'keyForEntity': 'dataJob', 'entityCategory': '_unset_', 'entityAspects': ['datahubIngestionRunSummary', 'datahubIngestionCheckpoint', 'domains', 'deprecation', 'versionInfo', 'dataJobInfo', 'dataJobInputOutput', 'editableDataJobProperties', 'ownership', 'status', 'globalTags', 'browsePaths', 'glossaryTerms', 'institutionalMemory', 'dataPlatformInstance', 'browsePathsV2', 'structuredProperties', 'forms', 'subTypes', 'incidentsSummary', 'testResults', 'dataTransformLogic']}
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.metadata.key.DataJobKey")
 
     def __init__(self,
@@ -24715,6 +24769,8 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.common.CostCostDiscriminator': CostCostDiscriminatorClass,
     'com.linkedin.pegasus2avro.common.CostType': CostTypeClass,
     'com.linkedin.pegasus2avro.common.DataPlatformInstance': DataPlatformInstanceClass,
+    'com.linkedin.pegasus2avro.common.DataTransform': DataTransformClass,
+    'com.linkedin.pegasus2avro.common.DataTransformLogic': DataTransformLogicClass,
     'com.linkedin.pegasus2avro.common.Deprecation': DeprecationClass,
     'com.linkedin.pegasus2avro.common.Documentation': DocumentationClass,
     'com.linkedin.pegasus2avro.common.DocumentationAssociation': DocumentationAssociationClass,
@@ -25182,6 +25238,8 @@ __SCHEMA_TYPES = {
     'CostCostDiscriminator': CostCostDiscriminatorClass,
     'CostType': CostTypeClass,
     'DataPlatformInstance': DataPlatformInstanceClass,
+    'DataTransform': DataTransformClass,
+    'DataTransformLogic': DataTransformLogicClass,
     'Deprecation': DeprecationClass,
     'Documentation': DocumentationClass,
     'DocumentationAssociation': DocumentationAssociationClass,
@@ -25588,6 +25646,7 @@ ASPECT_CLASSES: List[Type[_Aspect]] = [
     CostClass,
     BrowsePathsClass,
     InstitutionalMemoryClass,
+    DataTransformLogicClass,
     SubTypesClass,
     FormsClass,
     DeprecationClass,
@@ -25802,6 +25861,7 @@ class AspectBag(TypedDict, total=False):
     cost: CostClass
     browsePaths: BrowsePathsClass
     institutionalMemory: InstitutionalMemoryClass
+    dataTransformLogic: DataTransformLogicClass
     subTypes: SubTypesClass
     forms: FormsClass
     deprecation: DeprecationClass

datahub/metadata/com/linkedin/pegasus2avro/common/__init__.py

@@ -19,6 +19,8 @@ from .....schema_classes import CostCostClass
 from .....schema_classes import CostCostDiscriminatorClass
 from .....schema_classes import CostTypeClass
 from .....schema_classes import DataPlatformInstanceClass
+from .....schema_classes import DataTransformClass
+from .....schema_classes import DataTransformLogicClass
 from .....schema_classes import DeprecationClass
 from .....schema_classes import DocumentationClass
 from .....schema_classes import DocumentationAssociationClass
@@ -79,6 +81,8 @@ CostCost = CostCostClass
 CostCostDiscriminator = CostCostDiscriminatorClass
 CostType = CostTypeClass
 DataPlatformInstance = DataPlatformInstanceClass
+DataTransform = DataTransformClass
+DataTransformLogic = DataTransformLogicClass
 Deprecation = DeprecationClass
 Documentation = DocumentationClass
 DocumentationAssociation = DocumentationAssociationClass