acryl-datahub 1.3.0.1rc1__py3-none-any.whl → 1.3.0.1rc3__py3-none-any.whl

datahub/ingestion/source/unity/source.py

@@ -30,7 +30,6 @@ from datahub.emitter.mcp_builder import (
     add_entity_to_container,
     gen_containers,
 )
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -142,11 +141,7 @@ from datahub.metadata.schema_classes import (
 from datahub.metadata.urns import MlModelGroupUrn, MlModelUrn, TagUrn
 from datahub.sdk import MLModel, MLModelGroup
 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.utilities.file_backed_collections import FileBackedDict
 from datahub.utilities.hive_schema_to_avro import get_schema_fields_for_hive_column
 from datahub.utilities.registries.domain_registry import DomainRegistry
@@ -199,6 +194,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
     platform_resource_repository: Optional[UnityCatalogPlatformResourceRepository] = (
         None
     )
+    sql_parsing_aggregator: Optional[SqlParsingAggregator] = None
 
     def get_report(self) -> UnityCatalogReport:
         return self.report
@@ -218,6 +214,7 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             report=self.report,
             hive_metastore_proxy=self.hive_metastore_proxy,
             lineage_data_source=config.lineage_data_source,
+            usage_data_source=config.usage_data_source,
             databricks_api_page_size=config.databricks_api_page_size,
         )
 
@@ -244,9 +241,6 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
         self.table_refs: Set[TableReference] = set()
         self.view_refs: Set[TableReference] = set()
         self.notebooks: FileBackedDict[Notebook] = FileBackedDict()
-        self.view_definitions: FileBackedDict[Tuple[TableReference, str]] = (
-            FileBackedDict()
-        )
 
         # Global map of tables, for profiling
         self.tables: FileBackedDict[Table] = FileBackedDict()
@@ -290,6 +284,17 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                         platform_instance=self.config.platform_instance,
                         env=self.config.env,
                     )
+                    self.sql_parsing_aggregator = SqlParsingAggregator(
+                        platform=self.platform,
+                        platform_instance=self.config.platform_instance,
+                        env=self.config.env,
+                        schema_resolver=self.sql_parser_schema_resolver,
+                        generate_lineage=True,
+                        generate_queries=False,
+                        generate_usage_statistics=False,
+                        generate_operations=False,
+                    )
+                    self.report.sql_aggregator = self.sql_parsing_aggregator.report
             except Exception as e:
                 logger.debug("Exception", exc_info=True)
                 self.warn(
@@ -629,8 +634,13 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
             self.sql_parser_schema_resolver.add_schema_metadata(
                 dataset_urn, schema_metadata
             )
-            if table.view_definition:
-                self.view_definitions[dataset_urn] = (table.ref, table.view_definition)
+            if table.view_definition and self.sql_parsing_aggregator:
+                self.sql_parsing_aggregator.add_view_definition(
+                    view_urn=dataset_urn,
+                    view_definition=table.view_definition,
+                    default_db=table.ref.catalog,
+                    default_schema=table.ref.schema,
+                )
 
         if (
             table_props.customProperties.get("table_type")
@@ -1334,75 +1344,22 @@ class UnityCatalogSource(StatefulIngestionSourceBase, TestableSource):
                 )
             ]
 
-    def _run_sql_parser(
-        self, view_ref: TableReference, query: str, schema_resolver: SchemaResolver
-    ) -> Optional[SqlParsingResult]:
-        raw_lineage = sqlglot_lineage(
-            query,
-            schema_resolver=schema_resolver,
-            default_db=view_ref.catalog,
-            default_schema=view_ref.schema,
-        )
-        view_urn = self.gen_dataset_urn(view_ref)
-
-        if raw_lineage.debug_info.table_error:
-            logger.debug(
-                f"Failed to parse lineage for view {view_ref}: "
-                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_ref}: {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_ref}: {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)
-
     def get_view_lineage(self) -> Iterable[MetadataWorkUnit]:
         if not (
             self.config.include_hive_metastore
             and self.config.include_table_lineage
-            and self.sql_parser_schema_resolver
+            and self.sql_parsing_aggregator
         ):
             return
-        # This is only used for parsing view lineage. Usage, Operations are emitted elsewhere
-        builder = SqlParsingBuilder(
-            generate_lineage=True,
-            generate_usage_statistics=False,
-            generate_operations=False,
-        )
-        for dataset_name in self.view_definitions:
-            view_ref, view_definition = self.view_definitions[dataset_name]
-            result = self._run_sql_parser(
-                view_ref,
-                view_definition,
-                self.sql_parser_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,
-                )
-        yield from builder.gen_workunits()
+
+        for mcp in self.sql_parsing_aggregator.gen_metadata():
+            yield mcp.as_workunit()
 
     def close(self):
         if self.hive_metastore_proxy:
             self.hive_metastore_proxy.close()
-        if self.view_definitions:
-            self.view_definitions.close()
+        if self.sql_parsing_aggregator:
+            self.sql_parsing_aggregator.close()
         if self.sql_parser_schema_resolver:
             self.sql_parser_schema_resolver.close()
 

datahub/ingestion/source/unity/usage.py

@@ -11,7 +11,10 @@ from databricks.sdk.service.sql import QueryStatementType
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.source_helpers import auto_empty_dataset_usage_statistics
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.unity.config import UnityCatalogSourceConfig
+from datahub.ingestion.source.unity.config import (
+    UnityCatalogSourceConfig,
+    UsageDataSource,
+)
 from datahub.ingestion.source.unity.proxy import UnityCatalogApiProxy
 from datahub.ingestion.source.unity.proxy_types import (
     OPERATION_STATEMENT_TYPES,
@@ -164,11 +167,50 @@ class UnityCatalogUsageExtractor:
                     aspect=operation_aspect,
                 ).as_workunit()
 
+    def _validate_usage_data_source_config(self) -> None:
+        """Validate usage data source configuration before execution."""
+        usage_data_source = self.config.usage_data_source
+
+        if (
+            usage_data_source == UsageDataSource.SYSTEM_TABLES
+            and not self.proxy.warehouse_id
+        ):
+            raise ValueError(
+                "usage_data_source is set to SYSTEM_TABLES but warehouse_id is not configured. "
+                "Either set warehouse_id or use AUTO/API mode."
+            )
+
     def _get_queries(self) -> Iterable[Query]:
         try:
-            yield from self.proxy.query_history(
-                self.config.start_time, self.config.end_time
-            )
+            self._validate_usage_data_source_config()
+            usage_data_source = self.config.usage_data_source
+
+            if usage_data_source == UsageDataSource.AUTO:
+                if self.proxy.warehouse_id:
+                    logger.info(
+                        "Using system tables for usage query history (AUTO mode)"
+                    )
+                    yield from self.proxy.get_query_history_via_system_tables(
+                        self.config.start_time, self.config.end_time
+                    )
+                else:
+                    logger.info(
+                        "Using API for usage query history (AUTO mode, no warehouse)"
+                    )
+                    yield from self.proxy.query_history(
+                        self.config.start_time, self.config.end_time
+                    )
+            elif usage_data_source == UsageDataSource.SYSTEM_TABLES:
+                logger.info("Using system tables for usage query history (forced)")
+                yield from self.proxy.get_query_history_via_system_tables(
+                    self.config.start_time, self.config.end_time
+                )
+            elif usage_data_source == UsageDataSource.API:
+                logger.info("Using API for usage query history (forced)")
+                yield from self.proxy.query_history(
+                    self.config.start_time, self.config.end_time
+                )
+
         except Exception as e:
             logger.warning("Error getting queries", exc_info=True)
             self.report.report_warning("get-queries", str(e))

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -4,7 +4,6 @@ import enum
 import functools
 import json
 import logging
-import os
 import pathlib
 import tempfile
 import uuid
@@ -14,10 +13,10 @@ from typing import Callable, Dict, Iterable, List, Optional, Set, Union, cast
 
 import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
+from datahub.configuration.env_vars import get_sql_agg_query_log
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.emitter.mce_builder import get_sys_time, make_ts_millis
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.sql_parsing_builder import compute_upstream_fields
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.report import Report
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -84,7 +83,7 @@ class QueryLogSetting(enum.Enum):
 _DEFAULT_USER_URN = CorpUserUrn("_ingestion")
 _MISSING_SESSION_ID = "__MISSING_SESSION_ID"
 _DEFAULT_QUERY_LOG_SETTING = QueryLogSetting[
-    os.getenv("DATAHUB_SQL_AGG_QUERY_LOG") or QueryLogSetting.DISABLED.name
+    get_sql_agg_query_log() or QueryLogSetting.DISABLED.name
 ]
 MAX_UPSTREAM_TABLES_COUNT = 300
 MAX_FINEGRAINEDLINEAGE_COUNT = 2000
@@ -868,7 +867,7 @@ class SqlParsingAggregator(Closeable):
                 downstream=parsed.out_tables[0] if parsed.out_tables else None,
                 column_lineage=parsed.column_lineage,
                 # TODO: We need a full list of columns referenced, not just the out tables.
-                column_usage=compute_upstream_fields(parsed),
+                column_usage=self._compute_upstream_fields(parsed),
                 inferred_schema=infer_output_schema(parsed),
                 confidence_score=parsed.debug_info.confidence,
                 extra_info=observed.extra_info,
@@ -1157,7 +1156,7 @@ class SqlParsingAggregator(Closeable):
                 actor=None,
                 upstreams=parsed.in_tables,
                 column_lineage=parsed.column_lineage or [],
-                column_usage=compute_upstream_fields(parsed),
+                column_usage=self._compute_upstream_fields(parsed),
                 confidence_score=parsed.debug_info.confidence,
             )
         )
@@ -1741,6 +1740,16 @@ class SqlParsingAggregator(Closeable):
 
         return resolved_query
 
+    @staticmethod
+    def _compute_upstream_fields(
+        result: SqlParsingResult,
+    ) -> Dict[UrnStr, Set[UrnStr]]:
+        upstream_fields: Dict[UrnStr, Set[UrnStr]] = defaultdict(set)
+        for cl in result.column_lineage or []:
+            for upstream in cl.upstreams:
+                upstream_fields[upstream.table].add(upstream.column)
+        return upstream_fields
+
     def _gen_usage_statistics_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
         if not self._usage_aggregator:
             return

datahub/sql_parsing/sqlglot_lineage.py

@@ -691,6 +691,13 @@ def _column_level_lineage(
             select_statement=select_statement,
         )
 
+    # Handle VALUES expressions separately - they have no upstream tables and no column lineage
+    if isinstance(select_statement, sqlglot.exp.Values):
+        return _ColumnLineageWithDebugInfo(
+            column_lineage=[],
+            select_statement=select_statement,
+        )
+
     assert isinstance(select_statement, _SupportedColumnLineageTypesTuple)
     try:
         root_scope = sqlglot.optimizer.build_scope(select_statement)

datahub/telemetry/telemetry.py

@@ -16,6 +16,11 @@ from datahub._version import __version__, nice_version_name
 from datahub.cli.config_utils import DATAHUB_ROOT_FOLDER
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import ExceptionWithProps
+from datahub.configuration.env_vars import (
+    get_sentry_dsn,
+    get_sentry_environment,
+    get_telemetry_timeout,
+)
 from datahub.metadata.schema_classes import _custom_package_path
 from datahub.utilities.perf_timer import PerfTimer
 
@@ -97,11 +102,11 @@ if any(var in os.environ for var in CI_ENV_VARS):
 if _custom_package_path:
     ENV_ENABLED = False
 
-TIMEOUT = int(os.environ.get("DATAHUB_TELEMETRY_TIMEOUT", "10"))
+TIMEOUT = int(get_telemetry_timeout())
 MIXPANEL_ENDPOINT = "track.datahubproject.io/mp"
 MIXPANEL_TOKEN = "5ee83d940754d63cacbf7d34daa6f44a"
-SENTRY_DSN: Optional[str] = os.environ.get("SENTRY_DSN", None)
-SENTRY_ENVIRONMENT: str = os.environ.get("SENTRY_ENVIRONMENT", "dev")
+SENTRY_DSN: Optional[str] = get_sentry_dsn()
+SENTRY_ENVIRONMENT: str = get_sentry_environment()
 
 
 def _default_global_properties() -> Dict[str, Any]:

datahub/utilities/file_backed_collections.py

@@ -1,7 +1,6 @@
 import collections
 import gzip
 import logging
-import os
 import pathlib
 import pickle
 import shutil
@@ -28,6 +27,7 @@ from typing import (
     Union,
 )
 
+from datahub.configuration.env_vars import get_override_sqlite_version_req
 from datahub.ingestion.api.closeable import Closeable
 from datahub.utilities.sentinels import Unset, unset
 
@@ -36,7 +36,7 @@ logger: logging.Logger = logging.getLogger(__name__)
 
 def _get_sqlite_version_override() -> bool:
     """Check if SQLite version requirement should be overridden at runtime."""
-    override_str = os.environ.get("OVERRIDE_SQLITE_VERSION_REQ") or ""
+    override_str = get_override_sqlite_version_req()
     return bool(override_str and override_str.lower() != "false")
 
 

datahub/utilities/is_pytest.py

@@ -1,6 +1,7 @@
-import os
 import sys
 
+from datahub.configuration.env_vars import get_test_mode
+
 
 def is_pytest_running() -> bool:
-    return "pytest" in sys.modules and os.environ.get("DATAHUB_TEST_MODE") == "1"
+    return "pytest" in sys.modules and get_test_mode() == "1"

datahub/utilities/logging_manager.py

@@ -15,13 +15,13 @@ import collections
 import contextlib
 import itertools
 import logging
-import os
 import pathlib
 import sys
 from typing import Deque, Iterator, Optional
 
 import click
 
+from datahub.configuration.env_vars import get_no_color, get_suppress_logging_manager
 from datahub.utilities.tee_io import TeeIO
 
 BASE_LOGGING_FORMAT = (
@@ -38,7 +38,7 @@ IN_MEMORY_LOG_BUFFER_SIZE = 2000  # lines
 IN_MEMORY_LOG_BUFFER_MAX_LINE_LENGTH = 2000  # characters
 
 
-NO_COLOR = os.environ.get("NO_COLOR", False)
+NO_COLOR = get_no_color()
 
 
 def extract_name_from_filename(filename: str, fallback_name: str) -> str:
@@ -179,6 +179,18 @@ class _LogBuffer:
         return text
 
 
+class _ResilientStreamHandler(logging.StreamHandler):
+    """StreamHandler that gracefully handles closed streams."""
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            super().emit(record)
+        except (ValueError, OSError):
+            # Stream was closed (e.g., during pytest teardown)
+            # Silently ignore to prevent test failures
+            pass
+
+
 class _BufferLogHandler(logging.Handler):
     def __init__(self, storage: _LogBuffer) -> None:
         super().__init__()
@@ -201,7 +213,11 @@ class _BufferLogHandler(logging.Handler):
 def _remove_all_handlers(logger: logging.Logger) -> None:
     for handler in logger.handlers[:]:
         logger.removeHandler(handler)
-        handler.close()
+        try:
+            handler.close()
+        except (ValueError, OSError):
+            # Handler stream may already be closed (e.g., during pytest teardown)
+            pass
 
 
 _log_buffer = _LogBuffer(maxlen=IN_MEMORY_LOG_BUFFER_SIZE)
@@ -219,14 +235,14 @@ _default_formatter = logging.Formatter(BASE_LOGGING_FORMAT)
 def configure_logging(debug: bool, log_file: Optional[str] = None) -> Iterator[None]:
     _log_buffer.clear()
 
-    if os.environ.get("DATAHUB_SUPPRESS_LOGGING_MANAGER") == "1":
+    if get_suppress_logging_manager() == "1":
         # If we're running in pytest, we don't want to configure logging.
         yield
         return
 
     with contextlib.ExitStack() as stack:
         # Create stdout handler.
-        stream_handler = logging.StreamHandler()
+        stream_handler = _ResilientStreamHandler()
         stream_handler.addFilter(_DatahubLogFilter(debug=debug))
         stream_handler.setFormatter(_stream_formatter)
 
@@ -237,7 +253,7 @@ def configure_logging(debug: bool, log_file: Optional[str] = None) -> Iterator[N
             tee = TeeIO(sys.stdout, file)
             stack.enter_context(contextlib.redirect_stdout(tee))  # type: ignore
 
-            file_handler = logging.StreamHandler(file)
+            file_handler = _ResilientStreamHandler(file)
             file_handler.addFilter(_DatahubLogFilter(debug=True))
             file_handler.setFormatter(_default_formatter)
         else:

datahub/utilities/sample_data.py

@@ -1,12 +1,13 @@
-import os
 import pathlib
 import tempfile
 
 import requests
 
-DOCKER_COMPOSE_BASE = os.getenv(
-    "DOCKER_COMPOSE_BASE",
-    "https://raw.githubusercontent.com/datahub-project/datahub/master",
+from datahub.configuration.env_vars import get_docker_compose_base
+
+DOCKER_COMPOSE_BASE = (
+    get_docker_compose_base()
+    or "https://raw.githubusercontent.com/datahub-project/datahub/master"
 )
 BOOTSTRAP_MCES_FILE = "metadata-ingestion/examples/mce_files/bootstrap_mce.json"
 BOOTSTRAP_MCES_URL = f"{DOCKER_COMPOSE_BASE}/{BOOTSTRAP_MCES_FILE}"