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

datahub/ingestion/source/gc/dataprocess_cleanup.py

@@ -170,6 +170,8 @@ class DataProcessCleanupReport(SourceReport):
     sample_removed_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
@@ -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/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,

datahub/ingestion/source/sql/sql_common.py

@@ -11,7 +11,6 @@ from typing import (
     Dict,
     Iterable,
     List,
-    MutableMapping,
     Optional,
     Set,
     Tuple,
@@ -36,7 +35,6 @@ from datahub.emitter.mce_builder import (
     make_tag_urn,
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import capability
 from datahub.ingestion.api.incremental_lineage_helper import auto_incremental_lineage
@@ -79,7 +77,6 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
 from datahub.metadata.com.linkedin.pegasus2avro.common import StatusClass
-from datahub.metadata.com.linkedin.pegasus2avro.dataset import UpstreamLineage
 from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import DatasetSnapshot
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
@@ -106,17 +103,11 @@ from datahub.metadata.schema_classes import (
     GlobalTagsClass,
     SubTypesClass,
     TagAssociationClass,
-    UpstreamClass,
     ViewPropertiesClass,
 )
 from datahub.sql_parsing.schema_resolver import SchemaResolver
-from datahub.sql_parsing.sqlglot_lineage import (
-    SqlParsingResult,
-    sqlglot_lineage,
-    view_definition_lineage_helper,
-)
+from datahub.sql_parsing.sql_parsing_aggregator import SqlParsingAggregator
 from datahub.telemetry import telemetry
-from datahub.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.registries.domain_registry import DomainRegistry
 from datahub.utilities.sqlalchemy_type_converter import (
     get_native_data_type_for_sqlalchemy_type,
@@ -347,17 +338,19 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             )
 
         self.views_failed_parsing: Set[str] = set()
-        self.schema_resolver: SchemaResolver = SchemaResolver(
+
+        self.discovered_datasets: Set[str] = set()
+        self.aggregator = SqlParsingAggregator(
             platform=self.platform,
             platform_instance=self.config.platform_instance,
             env=self.config.env,
+            graph=self.ctx.graph,
+            generate_lineage=self.include_lineage,
+            generate_usage_statistics=False,
+            generate_operations=False,
+            eager_graph_load=False,
         )
-        self.discovered_datasets: Set[str] = set()
-        self._view_definition_cache: MutableMapping[str, str]
-        if self.config.use_file_backed_cache:
-            self._view_definition_cache = FileBackedDict[str]()
-        else:
-            self._view_definition_cache = {}
+        self.report.sql_aggregator = self.aggregator.report
 
     @classmethod
     def test_connection(cls, config_dict: dict) -> TestConnectionReport:
@@ -572,36 +565,9 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
                     profile_requests, profiler, platform=self.platform
                 )
 
-        if self.config.include_view_lineage:
-            yield from self.get_view_lineage()
-
-    def get_view_lineage(self) -> Iterable[MetadataWorkUnit]:
-        builder = SqlParsingBuilder(
-            generate_lineage=True,
-            generate_usage_statistics=False,
-            generate_operations=False,
-        )
-        for dataset_name in self._view_definition_cache.keys():
-            # TODO: Ensure that the lineage generated from the view definition
-            # matches the dataset_name.
-            view_definition = self._view_definition_cache[dataset_name]
-            result = self._run_sql_parser(
-                dataset_name,
-                view_definition,
-                self.schema_resolver,
-            )
-            if result and result.out_tables:
-                # This does not yield any workunits but we use
-                # yield here to execute this method
-                yield from builder.process_sql_parsing_result(
-                    result=result,
-                    query=view_definition,
-                    is_view_ddl=True,
-                    include_column_lineage=self.config.include_view_column_lineage,
-                )
-            else:
-                self.views_failed_parsing.add(dataset_name)
-        yield from builder.gen_workunits()
+        # Generate workunit for aggregated SQL parsing results
+        for mcp in self.aggregator.gen_metadata():
+            yield mcp.as_workunit()
 
     def get_identifier(
         self, *, schema: str, entity: str, inspector: Inspector, **kwargs: Any
@@ -760,16 +726,6 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
         )
         dataset_snapshot.aspects.append(dataset_properties)
 
-        if self.config.include_table_location_lineage and location_urn:
-            external_upstream_table = UpstreamClass(
-                dataset=location_urn,
-                type=DatasetLineageTypeClass.COPY,
-            )
-            yield MetadataChangeProposalWrapper(
-                entityUrn=dataset_snapshot.urn,
-                aspect=UpstreamLineage(upstreams=[external_upstream_table]),
-            ).as_workunit()
-
         extra_tags = self.get_extra_tags(inspector, schema, table)
         pk_constraints: dict = inspector.get_pk_constraint(table, schema)
         partitions: Optional[List[str]] = self.get_partitions(inspector, schema, table)
@@ -795,7 +751,7 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
 
         dataset_snapshot.aspects.append(schema_metadata)
         if self._save_schema_to_resolver():
-            self.schema_resolver.add_schema_metadata(dataset_urn, schema_metadata)
+            self.aggregator.register_schema(dataset_urn, schema_metadata)
             self.discovered_datasets.add(dataset_name)
         db_name = self.get_db_name(inspector)
 
@@ -815,6 +771,13 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             ),
         )
 
+        if self.config.include_table_location_lineage and location_urn:
+            self.aggregator.add_known_lineage_mapping(
+                upstream_urn=location_urn,
+                downstream_urn=dataset_snapshot.urn,
+                lineage_type=DatasetLineageTypeClass.COPY,
+            )
+
         if self.config.domain:
             assert self.domain_registry
             yield from get_domain_wu(
@@ -1089,6 +1052,7 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             self.config.platform_instance,
             self.config.env,
         )
+
         try:
             columns = inspector.get_columns(view, schema)
         except KeyError:
@@ -1108,7 +1072,7 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
                 canonical_schema=schema_fields,
             )
             if self._save_schema_to_resolver():
-                self.schema_resolver.add_schema_metadata(dataset_urn, schema_metadata)
+                self.aggregator.register_schema(dataset_urn, schema_metadata)
                 self.discovered_datasets.add(dataset_name)
 
         description, properties, _ = self.get_table_properties(inspector, schema, view)
@@ -1117,7 +1081,18 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
         view_definition = self._get_view_definition(inspector, schema, view)
         properties["view_definition"] = view_definition
         if view_definition and self.config.include_view_lineage:
-            self._view_definition_cache[dataset_name] = view_definition
+            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,
+            )
 
         dataset_snapshot = DatasetSnapshot(
             urn=dataset_urn,
@@ -1169,48 +1144,9 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             hasattr(self.config, "include_lineage") and self.config.include_lineage
         )
 
-    def _run_sql_parser(
-        self, view_identifier: str, query: str, schema_resolver: SchemaResolver
-    ) -> Optional[SqlParsingResult]:
-        try:
-            database, schema = self.get_db_schema(view_identifier)
-        except ValueError:
-            logger.warning(f"Invalid view identifier: {view_identifier}")
-            return None
-        raw_lineage = sqlglot_lineage(
-            query,
-            schema_resolver=schema_resolver,
-            default_db=database,
-            default_schema=schema,
-        )
-        view_urn = make_dataset_urn_with_platform_instance(
-            self.platform,
-            view_identifier,
-            self.config.platform_instance,
-            self.config.env,
-        )
-
-        if raw_lineage.debug_info.table_error:
-            logger.debug(
-                f"Failed to parse lineage for view {view_identifier}: "
-                f"{raw_lineage.debug_info.table_error}"
-            )
-            self.report.num_view_definitions_failed_parsing += 1
-            self.report.view_definitions_parsing_failures.append(
-                f"Table-level sql parsing error for view {view_identifier}: {raw_lineage.debug_info.table_error}"
-            )
-            return None
-
-        elif raw_lineage.debug_info.column_error:
-            self.report.num_view_definitions_failed_column_parsing += 1
-            self.report.view_definitions_parsing_failures.append(
-                f"Column-level sql parsing error for view {view_identifier}: {raw_lineage.debug_info.column_error}"
-            )
-        else:
-            self.report.num_view_definitions_parsed += 1
-            if raw_lineage.out_tables != [view_urn]:
-                self.report.num_view_definitions_view_urn_mismatch += 1
-        return view_definition_lineage_helper(raw_lineage, view_urn)
+    @property
+    def include_lineage(self):
+        return self.config.include_view_lineage
 
     def get_db_schema(self, dataset_identifier: str) -> Tuple[Optional[str], str]:
         database, schema, _view = dataset_identifier.split(".", 2)
@@ -1411,5 +1347,8 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             schema=schema, table=table, partition=partition, custom_sql=custom_sql
         )
 
+    def get_schema_resolver(self) -> SchemaResolver:
+        return self.aggregator._schema_resolver
+
     def get_report(self):
         return self.report

datahub/ingestion/source/sql/sql_generic_profiler.py

@@ -7,7 +7,10 @@ from typing import Dict, Iterable, List, Optional, Union, cast
 from sqlalchemy import create_engine, inspect
 from sqlalchemy.engine.reflection import Inspector
 
-from datahub.emitter.mce_builder import make_dataset_urn_with_platform_instance
+from datahub.emitter.mce_builder import (
+    make_dataset_urn_with_platform_instance,
+    parse_ts_millis,
+)
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.ge_data_profiler import (
@@ -245,11 +248,7 @@ class GenericProfiler:
                 # If profiling state exists we have to carry over to the new state
                 self.state_handler.add_to_state(dataset_urn, last_profiled)
 
-        threshold_time: Optional[datetime] = (
-            datetime.fromtimestamp(last_profiled / 1000, timezone.utc)
-            if last_profiled
-            else None
-        )
+        threshold_time: Optional[datetime] = parse_ts_millis(last_profiled)
         if (
             not threshold_time
             and self.config.profiling.profile_if_updated_since_days is not None

datahub/ingestion/source/sql/sql_report.py

@@ -5,6 +5,7 @@ from datahub.ingestion.glossary.classification_mixin import ClassificationReport
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
+from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.sqlalchemy_query_combiner import SQLAlchemyQueryCombinerReport
 from datahub.utilities.stats_collections import TopKDict, int_top_k_dict
@@ -52,6 +53,7 @@ class SQLSourceReport(
     num_view_definitions_failed_parsing: int = 0
     num_view_definitions_failed_column_parsing: int = 0
     view_definitions_parsing_failures: LossyList[str] = field(default_factory=LossyList)
+    sql_aggregator: Optional[SqlAggregatorReport] = None
 
     def report_entity_scanned(self, name: str, ent_type: str = "table") -> None:
         """

datahub/ingestion/source/state/checkpoint.py

@@ -12,6 +12,7 @@ from typing import Callable, Generic, Optional, Type, TypeVar
 import pydantic
 
 from datahub.configuration.common import ConfigModel
+from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.metadata.schema_classes import (
     DatahubIngestionCheckpointClass,
     IngestionCheckpointStateClass,
@@ -144,7 +145,7 @@ class Checkpoint(Generic[StateType]):
                 )
                 logger.info(
                     f"Successfully constructed last checkpoint state for job {job_name} "
-                    f"with timestamp {datetime.fromtimestamp(checkpoint_aspect.timestampMillis/1000, tz=timezone.utc)}"
+                    f"with timestamp {parse_ts_millis(checkpoint_aspect.timestampMillis)}"
                 )
                 return checkpoint
         return None

datahub/ingestion/source/tableau/tableau.py

@@ -920,10 +920,7 @@ class TableauSiteSource:
         return f"/{self.config.env.lower()}{self.no_env_browse_prefix}"
 
     def _re_authenticate(self) -> None:
-        self.report.info(
-            message="Re-authenticating to Tableau",
-            context=f"site='{self.site_content_url}'",
-        )
+        logger.info(f"Re-authenticating to Tableau site '{self.site_content_url}'")
         # Sign-in again may not be enough because Tableau sometimes caches invalid sessions
         # so we need to recreate the Tableau Server object
         self.server = self.config.make_tableau_client(self.site_content_url)

datahub/ingestion/source/unity/proxy.py

@@ -4,7 +4,7 @@ Manage the communication with DataBricks Server and provide equivalent dataclass
 
 import dataclasses
 import logging
-from datetime import datetime, timezone
+from datetime import datetime
 from typing import Any, Dict, Iterable, List, Optional, Union, cast
 from unittest.mock import patch
 
@@ -27,6 +27,7 @@ from databricks.sdk.service.sql import (
 from databricks.sdk.service.workspace import ObjectType
 
 import datahub
+from datahub.emitter.mce_builder import parse_ts_millis
 from datahub.ingestion.source.unity.hive_metastore_proxy import HiveMetastoreProxy
 from datahub.ingestion.source.unity.proxy_profiling import (
     UnityCatalogProxyProfilingMixin,
@@ -211,16 +212,8 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     id=obj.object_id,
                     path=obj.path,
                     language=obj.language,
-                    created_at=(
-                        datetime.fromtimestamp(obj.created_at / 1000, tz=timezone.utc)
-                        if obj.created_at
-                        else None
-                    ),
-                    modified_at=(
-                        datetime.fromtimestamp(obj.modified_at / 1000, tz=timezone.utc)
-                        if obj.modified_at
-                        else None
-                    ),
+                    created_at=parse_ts_millis(obj.created_at),
+                    modified_at=parse_ts_millis(obj.modified_at),
                 )
 
     def query_history(
@@ -452,17 +445,9 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             properties=obj.properties or {},
             owner=obj.owner,
             generation=obj.generation,
-            created_at=(
-                datetime.fromtimestamp(obj.created_at / 1000, tz=timezone.utc)
-                if obj.created_at
-                else None
-            ),
+            created_at=(parse_ts_millis(obj.created_at) if obj.created_at else None),
             created_by=obj.created_by,
-            updated_at=(
-                datetime.fromtimestamp(obj.updated_at / 1000, tz=timezone.utc)
-                if obj.updated_at
-                else None
-            ),
+            updated_at=(parse_ts_millis(obj.updated_at) if obj.updated_at else None),
             updated_by=obj.updated_by,
             table_id=obj.table_id,
             comment=obj.comment,
@@ -500,12 +485,8 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             query_id=info.query_id,
             query_text=info.query_text,
             statement_type=info.statement_type,
-            start_time=datetime.fromtimestamp(
-                info.query_start_time_ms / 1000, tz=timezone.utc
-            ),
-            end_time=datetime.fromtimestamp(
-                info.query_end_time_ms / 1000, tz=timezone.utc
-            ),
+            start_time=parse_ts_millis(info.query_start_time_ms),
+            end_time=parse_ts_millis(info.query_end_time_ms),
             user_id=info.user_id,
             user_name=info.user_name,
             executed_as_user_id=info.executed_as_user_id,