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

datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -1,7 +1,7 @@
 import logging
 from collections import defaultdict
 from dataclasses import dataclass, field
-from datetime import datetime, timezone
+from datetime import datetime
 from functools import lru_cache
 from typing import Any, Dict, FrozenSet, Iterable, Iterator, List, Optional
 
@@ -15,6 +15,7 @@ from google.cloud.bigquery.table import (
     TimePartitioningType,
 )
 
+from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.ingestion.api.source import SourceReport
 from datahub.ingestion.source.bigquery_v2.bigquery_audit import BigqueryTableIdentifier
 from datahub.ingestion.source.bigquery_v2.bigquery_helper import parse_labels
@@ -393,13 +394,7 @@ class BigQuerySchemaApi:
             name=table.table_name,
             created=table.created,
             table_type=table.table_type,
-            last_altered=(
-                datetime.fromtimestamp(
-                    table.get("last_altered") / 1000, tz=timezone.utc
-                )
-                if table.get("last_altered") is not None
-                else None
-            ),
+            last_altered=parse_ts_millis(table.get("last_altered")),
             size_in_bytes=table.get("bytes"),
             rows_count=table.get("row_count"),
             comment=table.comment,
@@ -460,11 +455,7 @@ class BigQuerySchemaApi:
         return BigqueryView(
             name=view.table_name,
             created=view.created,
-            last_altered=(
-                datetime.fromtimestamp(view.get("last_altered") / 1000, tz=timezone.utc)
-                if view.get("last_altered") is not None
-                else None
-            ),
+            last_altered=(parse_ts_millis(view.get("last_altered"))),
             comment=view.comment,
             view_definition=view.view_definition,
             materialized=view.table_type == BigqueryTableType.MATERIALIZED_VIEW,
@@ -705,13 +696,7 @@ class BigQuerySchemaApi:
         return BigqueryTableSnapshot(
             name=snapshot.table_name,
             created=snapshot.created,
-            last_altered=(
-                datetime.fromtimestamp(
-                    snapshot.get("last_altered") / 1000, tz=timezone.utc
-                )
-                if snapshot.get("last_altered") is not None
-                else None
-            ),
+            last_altered=parse_ts_millis(snapshot.get("last_altered")),
             comment=snapshot.comment,
             ddl=snapshot.ddl,
             snapshot_time=snapshot.snapshot_time,

datahub/ingestion/source/datahub/datahub_kafka_reader.py

@@ -12,6 +12,7 @@ from confluent_kafka.schema_registry import SchemaRegistryClient
 from confluent_kafka.schema_registry.avro import AvroDeserializer
 
 from datahub.configuration.kafka import KafkaConsumerConnectionConfig
+from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.source.datahub.config import DataHubSourceConfig
@@ -92,7 +93,7 @@ class DataHubKafkaReader(Closeable):
             if mcl.created and mcl.created.time > stop_time.timestamp() * 1000:
                 logger.info(
                     f"Stopped reading from kafka, reached MCL "
-                    f"with audit stamp {datetime.fromtimestamp(mcl.created.time / 1000)}"
+                    f"with audit stamp {parse_ts_millis(mcl.created.time)}"
                 )
                 break
 

datahub/ingestion/source/gc/dataprocess_cleanup.py

@@ -167,9 +167,11 @@ 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
+    num_data_jobs_found: int = 0
 
 
 class DataProcessCleanup:
@@ -265,13 +267,17 @@ class DataProcessCleanup:
                     self.report.report_failure(
                         f"Exception while deleting DPI: {e}", exc=e
                     )
-            if deleted_count_last_n % self.config.batch_size == 0:
+            if (
+                deleted_count_last_n % self.config.batch_size == 0
+                and deleted_count_last_n > 0
+            ):
                 logger.info(f"Deleted {deleted_count_last_n} DPIs from {job.urn}")
                 if self.config.delay:
                     logger.info(f"Sleeping for {self.config.delay} seconds")
                     time.sleep(self.config.delay)
 
-        logger.info(f"Deleted {deleted_count_last_n} DPIs from {job.urn}")
+        if deleted_count_last_n > 0:
+            logger.info(f"Deleted {deleted_count_last_n} DPIs from {job.urn}")
 
     def delete_entity(self, urn: str, type: str) -> None:
         assert self.ctx.graph
@@ -280,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(
@@ -351,7 +357,10 @@ class DataProcessCleanup:
             except Exception as e:
                 self.report.report_failure(f"Exception while deleting DPI: {e}", exc=e)
 
-            if deleted_count_retention % self.config.batch_size == 0:
+            if (
+                deleted_count_retention % self.config.batch_size == 0
+                and deleted_count_retention > 0
+            ):
                 logger.info(
                     f"Deleted {deleted_count_retention} DPIs from {job.urn} due to retention"
                 )
@@ -393,6 +402,7 @@ class DataProcessCleanup:
             scrollAcrossEntities = result.get("scrollAcrossEntities")
             if not scrollAcrossEntities:
                 raise ValueError("Missing scrollAcrossEntities in response")
+            self.report.num_data_flows_found += scrollAcrossEntities.get("count")
             logger.info(f"Got {scrollAcrossEntities.get('count')} DataFlow entities")
 
             scroll_id = scrollAcrossEntities.get("nextScrollId")
@@ -415,8 +425,9 @@ class DataProcessCleanup:
         assert self.ctx.graph
 
         dataFlows: Dict[str, DataFlowEntity] = {}
-        for flow in self.get_data_flows():
-            dataFlows[flow.urn] = flow
+        if self.config.delete_empty_data_flows:
+            for flow in self.get_data_flows():
+                dataFlows[flow.urn] = flow
 
         scroll_id: Optional[str] = None
         previous_scroll_id: Optional[str] = None
@@ -443,6 +454,7 @@ class DataProcessCleanup:
             if not scrollAcrossEntities:
                 raise ValueError("Missing scrollAcrossEntities in response")
 
+            self.report.num_data_jobs_found += scrollAcrossEntities.get("count")
             logger.info(f"Got {scrollAcrossEntities.get('count')} DataJob entities")
 
             scroll_id = scrollAcrossEntities.get("nextScrollId")
@@ -481,7 +493,8 @@ class DataProcessCleanup:
 
             previous_scroll_id = scroll_id
 
-        logger.info(f"Deleted {deleted_jobs} DataJobs")
+        if deleted_jobs > 0:
+            logger.info(f"Deleted {deleted_jobs} DataJobs")
         # Delete empty dataflows if needed
         if self.config.delete_empty_data_flows:
             deleted_data_flows: int = 0

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/ingestion/source/s3/source.py

@@ -225,7 +225,7 @@ class S3Source(StatefulIngestionSourceBase):
             self.init_spark()
 
     def init_spark(self):
-        os.environ.setdefault("SPARK_VERSION", "3.3")
+        os.environ.setdefault("SPARK_VERSION", "3.5")
         spark_version = os.environ["SPARK_VERSION"]
 
         # Importing here to avoid Deequ dependency for non profiling use cases

datahub/ingestion/source/sql/hive.py

@@ -838,3 +838,18 @@ class HiveSource(TwoTierSQLAlchemySource):
                 entityUrn=dataset_urn,
                 aspect=view_properties_aspect,
             ).as_workunit()
+
+        if view_definition and self.config.include_view_lineage:
+            default_db = None
+            default_schema = None
+            try:
+                default_db, default_schema = self.get_db_schema(dataset_name)
+            except ValueError:
+                logger.warning(f"Invalid view identifier: {dataset_name}")
+
+            self.aggregator.add_view_definition(
+                view_urn=dataset_urn,
+                view_definition=view_definition,
+                default_db=default_db,
+                default_schema=default_schema,
+            )

datahub/ingestion/source/sql/hive_metastore.py

@@ -123,6 +123,10 @@ class HiveMetastore(BasicSQLAlchemyConfig):
         description="Dataset Subtype name to be 'Table' or 'View' Valid options: ['True', 'False']",
     )
 
+    include_view_lineage: bool = Field(
+        default=False, description="", hidden_from_docs=True
+    )
+
     include_catalog_name_in_ids: bool = Field(
         default=False,
         description="Add the Presto catalog name (e.g. hive) to the generated dataset urns. `urn:li:dataset:(urn:li:dataPlatform:hive,hive.user.logging_events,PROD)` versus `urn:li:dataset:(urn:li:dataPlatform:hive,user.logging_events,PROD)`",
@@ -160,6 +164,9 @@ class HiveMetastore(BasicSQLAlchemyConfig):
 @capability(SourceCapability.DELETION_DETECTION, "Enabled via stateful ingestion")
 @capability(SourceCapability.DATA_PROFILING, "Not Supported", False)
 @capability(SourceCapability.CLASSIFICATION, "Not Supported", False)
+@capability(
+    SourceCapability.LINEAGE_COARSE, "View lineage is not supported", supported=False
+)
 class HiveMetastoreSource(SQLAlchemySource):
     """
     This plugin extracts the following:

datahub/ingestion/source/sql/mssql/source.py

@@ -724,7 +724,7 @@ class SQLServerSource(SQLAlchemySource):
             ):
                 yield from auto_workunit(
                     generate_procedure_lineage(
-                        schema_resolver=self.schema_resolver,
+                        schema_resolver=self.get_schema_resolver(),
                         procedure=procedure,
                         procedure_job_urn=MSSQLDataJob(entity=procedure).urn,
                         is_temp_table=self.is_temp_table,