acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/dremio/dremio_config.py

@@ -4,11 +4,12 @@ from typing import List, Literal, Optional
 import certifi
 from pydantic import Field, validator
 
-from datahub.configuration.common import AllowDenyPattern, ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
 )
+from datahub.configuration.time_window_config import BaseTimeWindowConfig
 from datahub.ingestion.source.ge_profiling_config import GEProfilingBaseConfig
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StatefulStaleMetadataRemovalConfig,
@@ -99,10 +100,9 @@ class ProfileConfig(GEProfilingBaseConfig):
     query_timeout: int = Field(
         default=300, description="Time before cancelling Dremio profiling query"
     )
-    include_field_median_value: bool = Field(
+    include_field_median_value: HiddenFromDocs[bool] = Field(
+        # Hidden because median causes a number of issues in Dremio.
         default=False,
-        hidden_from_docs=True,
-        description="Median causes a number of issues in Dremio.",
     )
 
 
@@ -118,6 +118,7 @@ class DremioSourceMapping(EnvConfigMixin, PlatformInstanceConfigMixin, ConfigMod
 class DremioSourceConfig(
     DremioConnectionConfig,
     StatefulIngestionConfigBase,
+    BaseTimeWindowConfig,
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
 ):

datahub/ingestion/source/dremio/dremio_reporting.py

@@ -1,22 +1,41 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime
+from typing import Optional
 
 from datahub.ingestion.source.sql.sql_report import SQLSourceReport
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
-from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
+from datahub.ingestion.source_report.time_window import BaseTimeWindowReport
+from datahub.sql_parsing.sql_parsing_aggregator import SqlAggregatorReport
+from datahub.utilities.stats_collections import (
+    TopKDict,
+    float_top_k_dict,
+    int_top_k_dict,
+)
 
 
 @dataclass
 class DremioSourceReport(
-    SQLSourceReport, StaleEntityRemovalSourceReport, IngestionStageReport
+    SQLSourceReport,
+    StaleEntityRemovalSourceReport,
+    BaseTimeWindowReport,
 ):
     num_containers_failed: int = 0
     num_datasets_failed: int = 0
     containers_scanned: int = 0
     containers_filtered: int = 0
 
+    api_calls_total: int = 0
+    api_calls_by_method_and_path: TopKDict[str, int] = field(
+        default_factory=int_top_k_dict
+    )
+    api_call_secs_by_method_and_path: TopKDict[str, float] = field(
+        default_factory=float_top_k_dict
+    )
+
+    sql_aggregator: Optional[SqlAggregatorReport] = None
+
     def report_upstream_latency(self, start_time: datetime, end_time: datetime) -> None:
         # recording total combined latency is not very useful, keeping this method as a placeholder
         # for future implementation of min / max / percentiles etc.

datahub/ingestion/source/dremio/dremio_source.py

@@ -22,6 +22,7 @@ from datahub.ingestion.api.source import (
     SourceReport,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.dremio.dremio_api import (
     DremioAPIOperations,
     DremioEdition,
@@ -51,13 +52,17 @@ 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,
+    IngestionHighStage,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
     DatasetLineageTypeClass,
     UpstreamClass,
     UpstreamLineage,
 )
-from datahub.metadata.schema_classes import ChangeTypeClass, SchemaMetadataClass
+from datahub.metadata.schema_classes import SchemaMetadataClass
 from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.sql_parsing_aggregator import (
     KnownQueryLineageInfo,
@@ -82,13 +87,34 @@ class DremioSourceMapEntry:
 @platform_name("Dremio")
 @config_class(DremioSourceConfig)
 @support_status(SupportStatus.CERTIFIED)
-@capability(SourceCapability.CONTAINERS, "Enabled by default")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.DREMIO_SPACE,
+        SourceCapabilityModifier.DREMIO_SOURCE,
+    ],
+)
 @capability(SourceCapability.DATA_PROFILING, "Optionally enabled via configuration")
 @capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
 @capability(SourceCapability.DOMAINS, "Supported via the `domain` config field")
-@capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
+@capability(
+    SourceCapability.LINEAGE_COARSE,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.TABLE,
+    ],
+)
+@capability(
+    SourceCapability.LINEAGE_FINE,
+    "Extract column-level lineage",
+    subtype_modifier=[
+        SourceCapabilityModifier.TABLE,
+    ],
+)
 @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 +152,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 +187,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 +224,88 @@ 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_high_stage(IngestionHighStage.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
@@ -300,10 +338,10 @@ class DremioSource(StatefulIngestionSourceBase):
             return
 
         dataset_urn = make_dataset_urn_with_platform_instance(
-            platform=make_data_platform_urn(self.get_platform()),
-            name=f"dremio.{dataset_name}",
-            env=self.config.env,
+            platform=self.get_platform(),
+            name=dataset_name,
             platform_instance=self.config.platform_instance,
+            env=self.config.env,
         )
 
         for dremio_mcp in self.dremio_aspects.populate_dataset_mcp(
@@ -383,13 +421,12 @@ class DremioSource(StatefulIngestionSourceBase):
         schema_str = ".".join(dataset_info.path)
         dataset_name = f"{schema_str}.{dataset_info.resource_name}".lower()
         dataset_urn = make_dataset_urn_with_platform_instance(
-            platform=make_data_platform_urn(self.get_platform()),
-            name=f"dremio.{dataset_name}",
-            env=self.config.env,
+            platform=self.get_platform(),
+            name=dataset_name,
             platform_instance=self.config.platform_instance,
+            env=self.config.env,
         )
-        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]
@@ -399,10 +436,10 @@ class DremioSource(StatefulIngestionSourceBase):
         """
         upstream_urns = [
             make_dataset_urn_with_platform_instance(
-                platform=make_data_platform_urn(self.get_platform()),
-                name=f"dremio.{upstream_table.lower()}",
-                env=self.config.env,
+                platform=self.get_platform(),
+                name=upstream_table.lower(),
                 platform_instance=self.config.platform_instance,
+                env=self.config.env,
             )
             for upstream_table in parents
         ]
@@ -417,11 +454,8 @@ class DremioSource(StatefulIngestionSourceBase):
             ]
         )
         mcp = MetadataChangeProposalWrapper(
-            entityType="dataset",
             entityUrn=dataset_urn,
-            aspectName=lineage.ASPECT_NAME,
             aspect=lineage,
-            changeType=ChangeTypeClass.UPSERT,
         )
 
         for upstream_urn in upstream_urns:
@@ -464,19 +498,19 @@ class DremioSource(StatefulIngestionSourceBase):
         if query.query and query.affected_dataset:
             upstream_urns = [
                 make_dataset_urn_with_platform_instance(
-                    platform=make_data_platform_urn(self.get_platform()),
-                    name=f"dremio.{ds.lower()}",
-                    env=self.config.env,
+                    platform=self.get_platform(),
+                    name=ds.lower(),
                     platform_instance=self.config.platform_instance,
+                    env=self.config.env,
                 )
                 for ds in query.queried_datasets
             ]
 
             downstream_urn = make_dataset_urn_with_platform_instance(
-                platform=make_data_platform_urn(self.get_platform()),
-                name=f"dremio.{query.affected_dataset.lower()}",
-                env=self.config.env,
+                platform=self.get_platform(),
+                name=query.affected_dataset.lower(),
                 platform_instance=self.config.platform_instance,
+                env=self.config.env,
             )
 
             # Add query to SqlParsingAggregator

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/dynamodb/dynamodb.py

@@ -12,7 +12,7 @@ from typing import (
     Union,
 )
 
-from pydantic.fields import Field
+from pydantic import Field, PositiveInt
 
 from datahub.configuration.common import AllowDenyPattern
 from datahub.configuration.source_common import DatasetSourceConfigMixin
@@ -73,7 +73,6 @@ from datahub.utilities.registries.domain_registry import DomainRegistry
 
 MAX_ITEMS_TO_RETRIEVE = 100
 PAGE_SIZE = 100
-MAX_SCHEMA_SIZE = 300
 MAX_PRIMARY_KEYS_SIZE = 100
 FIELD_DELIMITER = "."
 
@@ -107,6 +106,10 @@ class DynamoDBConfig(
         'Refer "Advanced Configurations" section for more details',
     )
 
+    max_schema_size: PositiveInt = Field(
+        default=300, description="Maximum number of fields to include in the schema."
+    )
+
     table_pattern: AllowDenyPattern = Field(
         default=AllowDenyPattern.allow_all(),
         description="Regex patterns for tables to filter in ingestion. The table name format is 'region.table'",
@@ -160,7 +163,7 @@ _attribute_type_to_field_type_mapping: Dict[str, Type] = {
 
 @platform_name("DynamoDB", id="dynamodb")
 @config_class(DynamoDBConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(
     SourceCapability.PLATFORM_INSTANCE,
     "By default, platform_instance will use the AWS account id",
@@ -455,25 +458,25 @@ class DynamoDBSource(StatefulIngestionSourceBase):
     ) -> SchemaMetadataClass:
         """ "
         To construct the schema metadata, it will first sort the schema by the occurrence of attribute names
-        in descending order and truncate the schema by MAX_SCHEMA_SIZE, and then start to construct the
+        in descending order and truncate the schema by max_schema_size, and then start to construct the
         schema metadata sorted by attribute name
         """
 
         canonical_schema: List[SchemaField] = []
         schema_size = len(schema.values())
         table_fields = list(schema.values())
-        if schema_size > MAX_SCHEMA_SIZE:
+        if schema_size > self.config.max_schema_size:
             # downsample the schema, using frequency as the sort key
             self.report.report_warning(
                 title="Schema Size Too Large",
-                message=f"Downsampling the table schema because MAX_SCHEMA_SIZE threshold is {MAX_SCHEMA_SIZE}",
+                message=f"Downsampling the table schema because `max_schema_size` threshold is {self.config.max_schema_size}",
                 context=f"Collection: {dataset_urn}",
             )
 
             # Add this information to the custom properties so user can know they are looking at down sampled schema
             dataset_properties.customProperties["schema.downsampled"] = "True"
             dataset_properties.customProperties["schema.totalFields"] = f"{schema_size}"
-        # append each schema field, schema will be sorted by count descending and delimited_name ascending and sliced to only include MAX_SCHEMA_SIZE items
+        # append each schema field, schema will be sorted by count descending and delimited_name ascending and sliced to only include max_schema_size items
         primary_keys = []
         for schema_field in sorted(
             table_fields,
@@ -481,7 +484,7 @@ class DynamoDBSource(StatefulIngestionSourceBase):
                 -x["count"],
                 x["delimited_name"],
             ),  # Negate `count` for descending order, `delimited_name` stays the same for ascending
-        )[0:MAX_SCHEMA_SIZE]:
+        )[: self.config.max_schema_size]:
             field_path = schema_field["delimited_name"]
             native_data_type = self.get_native_type(schema_field["type"], table_name)
             type = self.get_field_type(schema_field["type"], table_name)