acryl-datahub 0.15.0rc10__py3-none-any.whl → 0.15.0rc11__py3-none-any.whl

{acryl_datahub-0.15.0rc10.dist-info → acryl_datahub-0.15.0rc11.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datahub/__init__.py,sha256=vEQvZGPQ7hfoL7fm6XTBo6OY37wlMIITNNW1rpZuQCk,575
+datahub/__init__.py,sha256=RtQFMiAfUcKAAt_1AITPz1UOKRx2gIW0yLbLYUEChqU,575
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
 datahub/entrypoints.py,sha256=3-qSfXAx3Z0FEkBV5tlO8fQr4xk4ySeDRMVTpS5Xd6A,7793
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -61,7 +61,7 @@ datahub/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/cli/check_cli.py,sha256=9dXNyzZayHeoFjwFjLkMVyx6DiCZfeESyI-sYtGA6bE,12850
 datahub/cli/cli_utils.py,sha256=xMEK4tmoX2f_5BihxWrApOkayVLwhpPvMUU842x6FsI,13111
 datahub/cli/config_utils.py,sha256=yuXw7RzpRY5x_-MAoqWbv46qUkIeRNAJL4_OeJpYdBE,4879
-datahub/cli/delete_cli.py,sha256=YEbPDFcMtbIHDE0yNB9opfVI5r8f7kFOnfyc0rlyIqU,20906
+datahub/cli/delete_cli.py,sha256=Z7iXNr4ZMmghCnldU8laK4SwTNrhQEEnnUH_TeaBKog,21838
 datahub/cli/docker_check.py,sha256=rED4wHXqxcQ_qNFyIgFEZ85BHT9ZTE5YC-oUKqbRqi0,9432
 datahub/cli/docker_cli.py,sha256=QGoWFp8ZZsXOSMbgu0Q4snMmMmtP3epWAN-fYglUNEc,36491
 datahub/cli/env_utils.py,sha256=RQzjg4JE29hjPt4v7p-RuqoOr99w8E3DBHWiN2Sm7T4,252
@@ -97,7 +97,7 @@ datahub/configuration/git.py,sha256=s55eUHxKqVZgtVsISaDyS-1F4iZBiybbjYsjbp5LU5o,
 datahub/configuration/import_resolver.py,sha256=b4Ie9L7knN1LALEVMxTcNFSklDD6CVE-4Ipy4ZYhNYA,369
 datahub/configuration/json_loader.py,sha256=vIDnjwXWi9yHDO8KW64EupOzOb_sspehGCD7xGHzg84,302
 datahub/configuration/kafka.py,sha256=MlIwpd5FFyOyjdDXW_X9JTLNk7f988sPMgevkcZYVgI,2579
-datahub/configuration/kafka_consumer_config.py,sha256=14UWK6kTAnKTgHM43aeWjg67fddACsSLzhO6wgp6cm4,1175
+datahub/configuration/kafka_consumer_config.py,sha256=LivsObTt9yC3WoGnslJbF_x4ojfNdxMIMEhb8vvJfcA,2133
 datahub/configuration/pattern_utils.py,sha256=Q5IB9RfWOOo5FvRVBU7XkhiwHCxSQ1NTMfUlWtWI9qc,699
 datahub/configuration/pydantic_migration_helpers.py,sha256=4C_COAVZ5iJ8yxcWNgXZNWsY7ULogICNZ368oNF7zWg,1462
 datahub/configuration/source_common.py,sha256=68LZOuB23zSEcfgQJE1wZQnyYQHVVnEZK3Sniv_nEQs,2107
@@ -201,7 +201,7 @@ datahub/ingestion/source/nifi.py,sha256=ttsjZ9aRUvINmewvKFIQD8Rwa4jcl35WFG-F-jPG
 datahub/ingestion/source/openapi.py,sha256=3ea2ORz1cuq4e7L2hSjxG9Cw3__pVoJ5UNYTJS3EnKU,17386
 datahub/ingestion/source/openapi_parser.py,sha256=1_68wHWe_SzWYEyC1YVDw9vxoadKjW1yv8DecvyIhwY,13606
 datahub/ingestion/source/preset.py,sha256=eq7h1qKs8nfSBVot1ofN-YgZhw_rzq8DG4cKOGfDHko,3948
-datahub/ingestion/source/pulsar.py,sha256=bKtBzBTPuKwdTN4nSs7M0FYfA-q6FDc554W2NhnuWEY,19475
+datahub/ingestion/source/pulsar.py,sha256=H8XJC7xIX8Kdkd7006PxllAGVO_Pjza8Xx9VUBOvpPc,19827
 datahub/ingestion/source/redash.py,sha256=E-a14X19zppPun7_-S-pZ2lRiw1-68QiT-jL7bDzG10,32057
 datahub/ingestion/source/salesforce.py,sha256=S6LSM6mzl8-zKbrJPoINhM1SCpYfM244Xb74pbEI-J0,31792
 datahub/ingestion/source/source_registry.py,sha256=a2mLjJPLkSI-gYCTb_7U7Jo4D8jGknNQ_yScPIihXFk,1208
@@ -301,7 +301,7 @@ 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=f6Erj3KfD0Hx3ydwL5MUVCZgFzS9c6U2Pkr54JLIUOA,12394
-datahub/ingestion/source/gc/dataprocess_cleanup.py,sha256=rGIwUKPlNu2XR0YT8DuJGg2pPGIr3MM-YDa5Slo2vNY,14470
+datahub/ingestion/source/gc/dataprocess_cleanup.py,sha256=GOdLVYYfmmHV7Xtr-cwCTiQiRSVwq80Nu0EmARyoEeo,15323
 datahub/ingestion/source/gc/execution_request_cleanup.py,sha256=cHJmxz4NmA7VjTX2iGEo3wZ_SDrjC_rCQcnRxKgfUVI,8713
 datahub/ingestion/source/gc/soft_deleted_entity_cleanup.py,sha256=_tms5AqNAJRDRzQmyN_VydzXbdME2lkvTwa5u1La5z8,7353
 datahub/ingestion/source/gcs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -974,8 +974,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.0rc10.dist-info/METADATA,sha256=oHQM6bvFxv1QboFDXeahyZEbfacayYoFOlOUcPacaeI,172487
-acryl_datahub-0.15.0rc10.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-acryl_datahub-0.15.0rc10.dist-info/entry_points.txt,sha256=Yj0PWB0LQOq4Rj2fyR6ETx4BUGw4TOcNL0ZNoAZ9kQg,9504
-acryl_datahub-0.15.0rc10.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-0.15.0rc10.dist-info/RECORD,,
+acryl_datahub-0.15.0rc11.dist-info/METADATA,sha256=0GnLLvLM_Fvdw4XNusJLRrDZ0vemeu_B0ftz8XQXmbs,174408
+acryl_datahub-0.15.0rc11.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+acryl_datahub-0.15.0rc11.dist-info/entry_points.txt,sha256=Yj0PWB0LQOq4Rj2fyR6ETx4BUGw4TOcNL0ZNoAZ9kQg,9504
+acryl_datahub-0.15.0rc11.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-0.15.0rc11.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.0rc10"
+__version__ = "0.15.0rc11"
 
 
 def is_dev_mode() -> bool:

datahub/cli/delete_cli.py

@@ -214,14 +214,47 @@ def references(urn: str, dry_run: bool, force: bool) -> None:
 
 
 @delete.command()
-@click.option("--urn", required=True, type=str, help="the urn of the entity")
-def undo_by_filter(urn: str) -> None:
+@click.option("--urn", required=False, type=str, help="the urn of the entity")
+@click.option(
+    "-p",
+    "--platform",
+    required=False,
+    type=str,
+    help="Platform filter (e.g. snowflake)",
+)
+@click.option(
+    "-b",
+    "--batch-size",
+    required=False,
+    default=3000,
+    type=int,
+    help="Batch size when querying for entities to un-soft delete."
+    "Maximum 10000. Large batch sizes may cause timeouts.",
+)
+def undo_by_filter(
+    urn: Optional[str], platform: Optional[str], batch_size: int
+) -> None:
     """
-    Undo a soft deletion of an entity
+    Undo soft deletion by filters
     """
     graph = get_default_graph()
     logger.info(f"Using {graph}")
-    graph.set_soft_delete_status(urn=urn, delete=False)
+    if urn:
+        graph.set_soft_delete_status(urn=urn, delete=False)
+    else:
+        urns = list(
+            graph.get_urns_by_filter(
+                platform=platform,
+                query="*",
+                status=RemovedStatusFilter.ONLY_SOFT_DELETED,
+                batch_size=batch_size,
+            )
+        )
+        logger.info(f"Going to un-soft delete {len(urns)} urns")
+        urns_iter = progressbar.progressbar(urns, redirect_stdout=True)
+        for urn in urns_iter:
+            assert urn
+            graph.set_soft_delete_status(urn=urn, delete=False)
 
 
 @delete.command(no_args_is_help=True)

datahub/configuration/kafka_consumer_config.py

@@ -1,3 +1,4 @@
+import inspect
 import logging
 from typing import Any, Dict, Optional
 
@@ -34,5 +35,34 @@ class CallableConsumerConfig:
             "oauth_cb must be a string representing python function reference "
             "in the format <python-module>:<function-name>."
         )
+
+        call_back_fn = import_path(call_back)
+        self._validate_call_back_fn_signature(call_back_fn)
+
         # Set the callback
-        self._config[CallableConsumerConfig.CALLBACK_ATTRIBUTE] = import_path(call_back)
+        self._config[CallableConsumerConfig.CALLBACK_ATTRIBUTE] = call_back_fn
+
+    def _validate_call_back_fn_signature(self, call_back_fn: Any) -> None:
+        sig = inspect.signature(call_back_fn)
+
+        num_positional_args = len(
+            [
+                param
+                for param in sig.parameters.values()
+                if param.kind
+                in (
+                    inspect.Parameter.POSITIONAL_ONLY,
+                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                )
+                and param.default == inspect.Parameter.empty
+            ]
+        )
+
+        has_variadic_args = any(
+            param.kind == inspect.Parameter.VAR_POSITIONAL
+            for param in sig.parameters.values()
+        )
+
+        assert num_positional_args == 1 or (
+            has_variadic_args and num_positional_args <= 1
+        ), "oauth_cb function must accept single positional argument."

datahub/ingestion/source/gc/dataprocess_cleanup.py

@@ -208,22 +208,28 @@ class DataProcessCleanup:
         dpis = []
         start = 0
         while True:
-            job_query_result = self.ctx.graph.execute_graphql(
-                DATA_PROCESS_INSTANCES_QUERY,
-                {"dataJobUrn": job_urn, "start": start, "count": batch_size},
-            )
-            job_data = job_query_result.get("dataJob")
-            if not job_data:
-                raise ValueError(f"Error getting job {job_urn}")
-
-            runs_data = job_data.get("runs")
-            if not runs_data:
-                raise ValueError(f"Error getting runs for {job_urn}")
-
-            runs = runs_data.get("runs")
-            dpis.extend(runs)
-            start += batch_size
-            if len(runs) < batch_size:
+            try:
+                job_query_result = self.ctx.graph.execute_graphql(
+                    DATA_PROCESS_INSTANCES_QUERY,
+                    {"dataJobUrn": job_urn, "start": start, "count": batch_size},
+                )
+                job_data = job_query_result.get("dataJob")
+                if not job_data:
+                    logger.error(f"Error getting job {job_urn}")
+                    break
+
+                runs_data = job_data.get("runs")
+                if not runs_data:
+                    logger.error(f"Error getting runs for {job_urn}")
+                    break
+
+                runs = runs_data.get("runs")
+                dpis.extend(runs)
+                start += batch_size
+                if len(runs) < batch_size:
+                    break
+            except Exception as e:
+                logger.error(f"Exception while fetching DPIs for job {job_urn}: {e}")
                 break
         return dpis
 
@@ -243,8 +249,12 @@ class DataProcessCleanup:
                 futures[future] = dpi
 
             for future in as_completed(futures):
-                deleted_count_last_n += 1
-                futures[future]["deleted"] = True
+                try:
+                    future.result()
+                    deleted_count_last_n += 1
+                    futures[future]["deleted"] = True
+                except Exception as e:
+                    logger.error(f"Exception while deleting DPI: {e}")
 
             if deleted_count_last_n % self.config.batch_size == 0:
                 logger.info(f"Deleted {deleted_count_last_n} DPIs from {job.urn}")
@@ -279,7 +289,7 @@ class DataProcessCleanup:
         dpis = self.fetch_dpis(job.urn, self.config.batch_size)
         dpis.sort(
             key=lambda x: x["created"]["time"]
-            if x["created"] and x["created"]["time"]
+            if "created" in x and "time" in x["created"]
             else 0,
             reverse=True,
         )
@@ -314,15 +324,23 @@ class DataProcessCleanup:
             if dpi.get("deleted"):
                 continue
 
-            if dpi["created"]["time"] < retention_time * 1000:
+            if (
+                "created" not in dpi
+                or "time" not in dpi["created"]
+                or dpi["created"]["time"] < retention_time * 1000
+            ):
                 future = executor.submit(
                     self.delete_entity, dpi["urn"], "dataprocessInstance"
                 )
                 futures[future] = dpi
 
         for future in as_completed(futures):
-            deleted_count_retention += 1
-            futures[future]["deleted"] = True
+            try:
+                future.result()
+                deleted_count_retention += 1
+                futures[future]["deleted"] = True
+            except Exception as e:
+                logger.error(f"Exception while deleting DPI: {e}")
 
             if deleted_count_retention % self.config.batch_size == 0:
                 logger.info(
@@ -378,8 +396,11 @@ class DataProcessCleanup:
             dataFlows[flow.urn] = flow
 
         scroll_id: Optional[str] = None
+        previous_scroll_id: Optional[str] = None
+
         dataJobs: Dict[str, List[DataJobEntity]] = defaultdict(list)
         deleted_jobs: int = 0
+
         while True:
             result = self.ctx.graph.execute_graphql(
                 DATAJOB_QUERY,
@@ -426,9 +447,11 @@ class DataProcessCleanup:
                 else:
                     dataJobs[datajob_entity.flow_urn].append(datajob_entity)
 
-            if not scroll_id:
+            if not scroll_id or previous_scroll_id == scroll_id:
                 break
 
+            previous_scroll_id = scroll_id
+
         logger.info(f"Deleted {deleted_jobs} DataJobs")
         # Delete empty dataflows if needed
         if self.config.delete_empty_data_flows:
@@ -443,4 +466,5 @@ class DataProcessCleanup:
                     if deleted_jobs % self.config.batch_size == 0:
                         logger.info(f"Deleted {deleted_data_flows} DataFlows")
             logger.info(f"Deleted {deleted_data_flows} DataFlows")
+
         return []

datahub/ingestion/source/pulsar.py

@@ -78,7 +78,17 @@ class PulsarSchema:
     def __init__(self, schema):
         self.schema_version = schema.get("version")
 
-        avro_schema = json.loads(schema.get("data"))
+        schema_data = schema.get("data")
+        if not schema_data:
+            logger.warning("Schema data is empty or None. Using default empty schema.")
+            schema_data = "{}"
+
+        try:
+            avro_schema = json.loads(schema_data)
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON schema: {schema_data}. Error: {str(e)}")
+            avro_schema = {}
+
         self.schema_name = avro_schema.get("namespace") + "." + avro_schema.get("name")
         self.schema_description = avro_schema.get("doc")
         self.schema_type = schema.get("type")