acryl-datahub 1.1.0.4rc2__py3-none-any.whl → 1.1.0.4rc3__py3-none-any.whl

datahub/ingestion/source/dremio/dremio_source.py

@@ -51,7 +51,11 @@ from datahub.ingestion.source.state.stale_entity_removal_handler import (
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionSourceBase,
 )
-from datahub.ingestion.source_report.ingestion_stage import PROFILING
+from datahub.ingestion.source_report.ingestion_stage import (
+    LINEAGE_EXTRACTION,
+    METADATA_EXTRACTION,
+    PROFILING,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     DatasetLineageTypeClass,
     UpstreamClass,
@@ -89,6 +93,7 @@ class DremioSourceMapEntry:
 @capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
 @capability(SourceCapability.OWNERSHIP, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 class DremioSource(StatefulIngestionSourceBase):
     """
     This plugin integrates with Dremio to extract and ingest metadata into DataHub.
@@ -126,6 +131,13 @@ class DremioSource(StatefulIngestionSourceBase):
         self.default_db = "dremio"
         self.config = config
         self.report = DremioSourceReport()
+
+        # Set time window for query lineage extraction
+        self.report.window_start_time, self.report.window_end_time = (
+            self.config.start_time,
+            self.config.end_time,
+        )
+
         self.source_map: Dict[str, DremioSourceMapEntry] = dict()
 
         # Initialize API operations
@@ -154,6 +166,7 @@ class DremioSource(StatefulIngestionSourceBase):
             generate_operations=True,
             usage_config=self.config.usage,
         )
+        self.report.sql_aggregator = self.sql_parsing_aggregator.report
 
         # For profiling
         self.profiler = DremioProfiler(config, self.report, dremio_api)
@@ -190,84 +203,85 @@ class DremioSource(StatefulIngestionSourceBase):
 
         self.source_map = self._build_source_map()
 
-        # Process Containers
-        containers = self.dremio_catalog.get_containers()
-        for container in containers:
-            try:
-                yield from self.process_container(container)
-                logger.info(
-                    f"Dremio container {container.container_name} emitted successfully"
-                )
-            except Exception as exc:
-                self.report.num_containers_failed += 1  # Increment failed containers
-                self.report.report_failure(
-                    message="Failed to process Dremio container",
-                    context=f"{'.'.join(container.path)}.{container.container_name}",
-                    exc=exc,
-                )
+        with self.report.new_stage(METADATA_EXTRACTION):
+            # Process Containers
+            containers = self.dremio_catalog.get_containers()
+            for container in containers:
+                try:
+                    yield from self.process_container(container)
+                    logger.info(
+                        f"Dremio container {container.container_name} emitted successfully"
+                    )
+                except Exception as exc:
+                    self.report.num_containers_failed += 1
+                    self.report.report_failure(
+                        message="Failed to process Dremio container",
+                        context=f"{'.'.join(container.path)}.{container.container_name}",
+                        exc=exc,
+                    )
 
-        # Process Datasets
-        datasets = self.dremio_catalog.get_datasets()
+            # Process Datasets
+            datasets = self.dremio_catalog.get_datasets()
 
-        for dataset_info in datasets:
-            try:
-                yield from self.process_dataset(dataset_info)
-                logger.info(
-                    f"Dremio dataset {'.'.join(dataset_info.path)}.{dataset_info.resource_name} emitted successfully"
-                )
-            except Exception as exc:
-                self.report.num_datasets_failed += 1  # Increment failed datasets
-                self.report.report_failure(
-                    message="Failed to process Dremio dataset",
-                    context=f"{'.'.join(dataset_info.path)}.{dataset_info.resource_name}",
-                    exc=exc,
-                )
+            for dataset_info in datasets:
+                try:
+                    yield from self.process_dataset(dataset_info)
+                    logger.info(
+                        f"Dremio dataset {'.'.join(dataset_info.path)}.{dataset_info.resource_name} emitted successfully"
+                    )
+                except Exception as exc:
+                    self.report.num_datasets_failed += 1  # Increment failed datasets
+                    self.report.report_failure(
+                        message="Failed to process Dremio dataset",
+                        context=f"{'.'.join(dataset_info.path)}.{dataset_info.resource_name}",
+                        exc=exc,
+                    )
 
-        # Optionally Process Query Lineage
-        if self.config.include_query_lineage:
-            self.get_query_lineage_workunits()
-
-        # Process Glossary Terms
-        glossary_terms = self.dremio_catalog.get_glossary_terms()
-
-        for glossary_term in glossary_terms:
-            try:
-                yield from self.process_glossary_term(glossary_term)
-            except Exception as exc:
-                self.report.report_failure(
-                    message="Failed to process Glossary terms",
-                    context=f"{glossary_term.glossary_term}",
-                    exc=exc,
-                )
+            # Process Glossary Terms
+            glossary_terms = self.dremio_catalog.get_glossary_terms()
 
-        # Generate workunit for aggregated SQL parsing results
-        for mcp in self.sql_parsing_aggregator.gen_metadata():
-            self.report.report_workunit(mcp.as_workunit())
-            yield mcp.as_workunit()
-
-        # Profiling
-        if self.config.is_profiling_enabled():
-            with ThreadPoolExecutor(
-                max_workers=self.config.profiling.max_workers
-            ) as executor:
-                future_to_dataset = {
-                    executor.submit(self.generate_profiles, dataset): dataset
-                    for dataset in datasets
-                }
-
-                for future in as_completed(future_to_dataset):
-                    dataset_info = future_to_dataset[future]
-                    try:
-                        yield from future.result()
-                    except Exception as exc:
-                        self.report.profiling_skipped_other[
-                            dataset_info.resource_name
-                        ] += 1
-                        self.report.report_failure(
-                            message="Failed to profile dataset",
-                            context=f"{'.'.join(dataset_info.path)}.{dataset_info.resource_name}",
-                            exc=exc,
-                        )
+            for glossary_term in glossary_terms:
+                try:
+                    yield from self.process_glossary_term(glossary_term)
+                except Exception as exc:
+                    self.report.report_failure(
+                        message="Failed to process Glossary terms",
+                        context=f"{glossary_term.glossary_term}",
+                        exc=exc,
+                    )
+
+            # Optionally Process Query Lineage
+            if self.config.include_query_lineage:
+                with self.report.new_stage(LINEAGE_EXTRACTION):
+                    self.get_query_lineage_workunits()
+
+            # Generate workunit for aggregated SQL parsing results
+            for mcp in self.sql_parsing_aggregator.gen_metadata():
+                yield mcp.as_workunit()
+
+            # Profiling
+            if self.config.is_profiling_enabled():
+                with self.report.new_stage(PROFILING), ThreadPoolExecutor(
+                    max_workers=self.config.profiling.max_workers
+                ) as executor:
+                    future_to_dataset = {
+                        executor.submit(self.generate_profiles, dataset): dataset
+                        for dataset in datasets
+                    }
+
+                    for future in as_completed(future_to_dataset):
+                        dataset_info = future_to_dataset[future]
+                        try:
+                            yield from future.result()
+                        except Exception as exc:
+                            self.report.profiling_skipped_other[
+                                dataset_info.resource_name
+                            ] += 1
+                            self.report.report_failure(
+                                message="Failed to profile dataset",
+                                context=f"{'.'.join(dataset_info.path)}.{dataset_info.resource_name}",
+                                exc=exc,
+                            )
 
     def process_container(
         self, container_info: DremioContainer
@@ -388,8 +402,7 @@ class DremioSource(StatefulIngestionSourceBase):
             env=self.config.env,
             platform_instance=self.config.platform_instance,
         )
-        with self.report.new_stage(f"{dataset_info.resource_name}: {PROFILING}"):
-            yield from self.profiler.get_workunits(dataset_info, dataset_urn)
+        yield from self.profiler.get_workunits(dataset_info, dataset_urn)
 
     def generate_view_lineage(
         self, dataset_urn: str, parents: List[str]

datahub/ingestion/source/dremio/dremio_sql_queries.py

@@ -1,3 +1,7 @@
+from datetime import datetime, timedelta
+from typing import Optional
+
+
 class DremioSQLQueries:
     QUERY_DATASETS_CE = """
     SELECT* FROM
@@ -235,28 +239,83 @@ class DremioSQLQueries:
             TABLE_NAME ASC
             """
 
-    # Dremio Documentation: https://docs.dremio.com/current/reference/sql/system-tables/jobs_recent/
-    # queried_datasets incorrectly documented as [varchar]. Observed as varchar.
-    # LENGTH used as opposed to ARRAY_SIZE
-    QUERY_ALL_JOBS = """
-    SELECT
-        job_id,
-        user_name,
-        submitted_ts,
-        query,
-        queried_datasets
-    FROM
-        SYS.JOBS_RECENT
-    WHERE
-        STATUS = 'COMPLETED'
-        AND LENGTH(queried_datasets)>0
-        AND user_name != '$dremio$'
-        AND query_type not like '%INTERNAL%'
-    """
+    @staticmethod
+    def _get_default_start_timestamp_millis() -> str:
+        """Get default start timestamp (1 day ago) in milliseconds precision format"""
+        one_day_ago = datetime.now() - timedelta(days=1)
+        return one_day_ago.strftime("%Y-%m-%d %H:%M:%S.%f")[
+            :-3
+        ]  # Truncate to milliseconds
+
+    @staticmethod
+    def _get_default_end_timestamp_millis() -> str:
+        """Get default end timestamp (now) in milliseconds precision format"""
+        now = datetime.now()
+        return now.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]  # Truncate to milliseconds
+
+    @staticmethod
+    def get_query_all_jobs(
+        start_timestamp_millis: Optional[str] = None,
+        end_timestamp_millis: Optional[str] = None,
+    ) -> str:
+        """
+        Get query for all jobs with optional time filtering.
+
+        Args:
+            start_timestamp_millis: Start timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to 1 day ago)
+            end_timestamp_millis: End timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to now)
+
+        Returns:
+            SQL query string with time filtering applied
+        """
+        if start_timestamp_millis is None:
+            start_timestamp_millis = (
+                DremioSQLQueries._get_default_start_timestamp_millis()
+            )
+        if end_timestamp_millis is None:
+            end_timestamp_millis = DremioSQLQueries._get_default_end_timestamp_millis()
+
+        return f"""
+        SELECT
+            job_id,
+            user_name,
+            submitted_ts,
+            query,
+            queried_datasets
+        FROM
+            SYS.JOBS_RECENT
+        WHERE
+            STATUS = 'COMPLETED'
+            AND LENGTH(queried_datasets)>0
+            AND user_name != '$dremio$'
+            AND query_type not like '%INTERNAL%'
+            AND submitted_ts >= TIMESTAMP '{start_timestamp_millis}'
+            AND submitted_ts <= TIMESTAMP '{end_timestamp_millis}'
+        """
+
+    @staticmethod
+    def get_query_all_jobs_cloud(
+        start_timestamp_millis: Optional[str] = None,
+        end_timestamp_millis: Optional[str] = None,
+    ) -> str:
+        """
+        Get query for all jobs in Dremio Cloud with optional time filtering.
+
+        Args:
+            start_timestamp_millis: Start timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to 7 days ago)
+            end_timestamp_millis: End timestamp in format 'YYYY-MM-DD HH:MM:SS.mmm' (defaults to now)
+
+        Returns:
+            SQL query string with time filtering applied
+        """
+        if start_timestamp_millis is None:
+            start_timestamp_millis = (
+                DremioSQLQueries._get_default_start_timestamp_millis()
+            )
+        if end_timestamp_millis is None:
+            end_timestamp_millis = DremioSQLQueries._get_default_end_timestamp_millis()
 
-    # Dremio Documentation: https://docs.dremio.com/cloud/reference/sql/system-tables/jobs-historical
-    # queried_datasets correctly documented as [varchar]
-    QUERY_ALL_JOBS_CLOUD = """
+        return f"""
         SELECT
             job_id,
             user_name,
@@ -270,6 +329,8 @@ class DremioSQLQueries:
             AND ARRAY_SIZE(queried_datasets)>0
             AND user_name != '$dremio$'
             AND query_type not like '%INTERNAL%'
+            AND submitted_ts >= TIMESTAMP '{start_timestamp_millis}'
+            AND submitted_ts <= TIMESTAMP '{end_timestamp_millis}'
         """
 
     QUERY_TYPES = [

datahub/ingestion/source/ge_data_profiler.py

@@ -120,7 +120,6 @@ SNOWFLAKE = "snowflake"
 BIGQUERY = "bigquery"
 REDSHIFT = "redshift"
 DATABRICKS = "databricks"
-TRINO = "trino"
 
 # Type names for Databricks, to match Title Case types in sqlalchemy
 ProfilerTypeMapping.INT_TYPE_NAMES.append("Integer")
@@ -206,6 +205,17 @@ def get_column_unique_count_dh_patch(self: SqlAlchemyDataset, column: str) -> in
             )
         )
         return convert_to_json_serializable(element_values.fetchone()[0])
+    elif (
+        self.engine.dialect.name.lower() == GXSqlDialect.AWSATHENA
+        or self.engine.dialect.name.lower() == GXSqlDialect.TRINO
+    ):
+        return convert_to_json_serializable(
+            self.engine.execute(
+                sa.select(sa.func.approx_distinct(sa.column(column))).select_from(
+                    self._table
+                )
+            ).scalar()
+        )
     return convert_to_json_serializable(
         self.engine.execute(
             sa.select([sa.func.count(sa.func.distinct(sa.column(column)))]).select_from(
@@ -734,11 +744,41 @@ class _SingleDatasetProfiler(BasicDatasetProfilerBase):
     def _get_dataset_column_distinct_value_frequencies(
         self, column_profile: DatasetFieldProfileClass, column: str
     ) -> None:
-        if self.config.include_field_distinct_value_frequencies:
+        if not self.config.include_field_distinct_value_frequencies:
+            return
+        try:
+            results = self.dataset.engine.execute(
+                sa.select(
+                    [
+                        sa.column(column),
+                        sa.func.count(sa.column(column)),
+                    ]
+                )
+                .select_from(self.dataset._table)
+                .where(sa.column(column).is_not(None))
+                .group_by(sa.column(column))
+            ).fetchall()
+
             column_profile.distinctValueFrequencies = [
-                ValueFrequencyClass(value=str(value), frequency=count)
-                for value, count in self.dataset.get_column_value_counts(column).items()
+                ValueFrequencyClass(value=str(value), frequency=int(count))
+                for value, count in results
             ]
+            # sort so output is deterministic. don't do it in SQL because not all column
+            # types are sortable in SQL (such as JSON data types on Athena/Trino).
+            column_profile.distinctValueFrequencies = sorted(
+                column_profile.distinctValueFrequencies, key=lambda x: x.value
+            )
+        except Exception as e:
+            logger.debug(
+                f"Caught exception while attempting to get distinct value frequencies for column {column}. {e}"
+            )
+
+            self.report.report_warning(
+                title="Profiling: Unable to Calculate Distinct Value Frequencies",
+                message="Distinct value frequencies for the column will not be accessible",
+                context=f"{self.dataset_name}.{column}",
+                exc=e,
+            )
 
     @_run_with_query_combiner
     def _get_dataset_column_histogram(
@@ -1395,12 +1435,12 @@ class DatahubGEProfiler:
                     )
                 return None
             finally:
-                if batch is not None and self.base_engine.engine.name.upper() in [
-                    "TRINO",
-                    "AWSATHENA",
+                if batch is not None and self.base_engine.engine.name.lower() in [
+                    GXSqlDialect.TRINO,
+                    GXSqlDialect.AWSATHENA,
                 ]:
                     if (
-                        self.base_engine.engine.name.upper() == "TRINO"
+                        self.base_engine.engine.name.lower() == GXSqlDialect.TRINO
                         or temp_view is not None
                     ):
                         self._drop_temp_table(batch)

datahub/ingestion/source/redshift/redshift.py

@@ -10,6 +10,7 @@ import humanfriendly
 import pydantic
 import redshift_connector
 
+from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.pattern_utils import is_schema_allowed
 from datahub.emitter.mce_builder import (
     make_data_platform_urn,
@@ -357,7 +358,23 @@ class RedshiftSource(StatefulIngestionSourceBase, TestableSource):
             ).workunit_processor,
         ]
 
+    def _warn_deprecated_configs(self):
+        if (
+            self.config.match_fully_qualified_names is not None
+            and not self.config.match_fully_qualified_names
+            and self.config.schema_pattern is not None
+            and self.config.schema_pattern != AllowDenyPattern.allow_all()
+        ):
+            self.report.report_warning(
+                message="Please update `schema_pattern` to match against fully qualified schema name `<database_name>.<schema_name>` and set config `match_fully_qualified_names : True`."
+                "Current default `match_fully_qualified_names: False` is only to maintain backward compatibility. "
+                "The config option `match_fully_qualified_names` will be removed in future and the default behavior will be like `match_fully_qualified_names: True`.",
+                context="Config option deprecation warning",
+                title="Config option deprecation warning",
+            )
+
     def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        self._warn_deprecated_configs()
         connection = self._try_get_redshift_connection(self.config)
 
         if connection is None:

datahub/ingestion/source/usage/clickhouse_usage.py

@@ -89,6 +89,7 @@ class ClickHouseUsageConfig(ClickHouseConfig, BaseUsageConfig, EnvConfigMixin):
     SourceCapability.DELETION_DETECTION, "Enabled by default via stateful ingestion"
 )
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 @dataclasses.dataclass
 class ClickHouseUsageSource(Source):
     """

datahub/ingestion/source/usage/starburst_trino_usage.py

@@ -15,7 +15,9 @@ from sqlalchemy.engine import Engine
 import datahub.emitter.mce_builder as builder
 from datahub.configuration.time_window_config import get_time_bucket
 from datahub.ingestion.api.decorators import (
+    SourceCapability,
     SupportStatus,
+    capability,
     config_class,
     platform_name,
     support_status,
@@ -112,6 +114,7 @@ class TrinoUsageReport(SourceReport):
 @platform_name("Trino")
 @config_class(TrinoUsageConfig)
 @support_status(SupportStatus.CERTIFIED)
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 @dataclasses.dataclass
 class TrinoUsageSource(Source):
     """

datahub/utilities/stats_collections.py

@@ -56,3 +56,7 @@ class TopKDict(DefaultDict[_KT, _VT]):
 
 def int_top_k_dict() -> TopKDict[str, int]:
     return TopKDict(int)
+
+
+def float_top_k_dict() -> TopKDict[str, float]:
+    return TopKDict(float)