acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/dremio/dremio_source.py

@@ -1,7 +1,6 @@
 import logging
-import re
-from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
 from typing import Dict, Iterable, List, Optional
 
 from datahub.emitter.mce_builder import (
@@ -23,12 +22,16 @@ 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,
 )
 from datahub.ingestion.source.dremio.dremio_aspects import DremioAspects
-from datahub.ingestion.source.dremio.dremio_config import DremioSourceConfig
+from datahub.ingestion.source.dremio.dremio_config import (
+    DremioSourceConfig,
+    DremioSourceMapping,
+)
 from datahub.ingestion.source.dremio.dremio_datahub_source_mapping import (
     DremioToDataHubSourceTypeMapping,
 )
@@ -39,6 +42,7 @@ from datahub.ingestion.source.dremio.dremio_entities import (
     DremioDatasetType,
     DremioGlossaryTerm,
     DremioQuery,
+    DremioSourceContainer,
 )
 from datahub.ingestion.source.dremio.dremio_profiling import DremioProfiler
 from datahub.ingestion.source.dremio.dremio_reporting import DremioSourceReport
@@ -48,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,
@@ -65,16 +73,48 @@ from datahub.sql_parsing.sql_parsing_aggregator import (
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class DremioSourceMapEntry:
+    platform: str
+    source_name: str
+    dremio_source_category: str
+    root_path: str = ""
+    database_name: str = ""
+    platform_instance: Optional[str] = None
+    env: Optional[str] = None
+
+
 @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.
@@ -112,7 +152,14 @@ class DremioSource(StatefulIngestionSourceBase):
         self.default_db = "dremio"
         self.config = config
         self.report = DremioSourceReport()
-        self.source_map: Dict[str, Dict] = defaultdict()
+
+        # 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
         dremio_api = DremioAPIOperations(self.config, self.report)
@@ -140,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)
@@ -152,111 +200,12 @@ class DremioSource(StatefulIngestionSourceBase):
     def get_platform(self) -> str:
         return "dremio"
 
-    def _build_source_map(self) -> Dict[str, Dict]:
-        """
-        Builds a source mapping dictionary to support external lineage generation across
-        multiple Dremio sources, based on provided configuration mappings.
-
-        This method operates as follows:
-
-        1. If a source mapping is present in the config:
-        - For each source in the Dremio catalog, if the mapping's `source_name` matches
-            the `dremio_source_type`, `root_path` and `database_name` are added to the mapping
-            information, along with the platform, platform instance, and environment if they exist.
-            This allows constructing the full URN for upstream lineage.
-
-        2. If a source mapping is absent in the configuration:
-        - Default mappings are created for each source name, setting `env` and `platform_instance`
-            to default values and classifying the source type. This ensures all sources have a
-            mapping, even if specific configuration details are missing.
-
-        Returns:
-            Dict[str, Dict]: A dictionary (`source_map`) where each key is a source name
-                            (lowercased) and each value is another dictionary containing:
-                            - `platform`: The source platform.
-                            - `source_name`: The source name.
-                            - `dremio_source_type`: The type mapped to DataHub,
-                            e.g., "database", "folder".
-                            - Optional `root_path`, `database_name`, `platform_instance`,
-                            and `env` if provided in the configuration.
-        Example:
-            This method is used internally within the class to generate mappings before
-            creating cross-platform lineage.
-
-        """
-
-        source_map = {}
+    def _build_source_map(self) -> Dict[str, DremioSourceMapEntry]:
         dremio_sources = self.dremio_catalog.get_sources()
+        source_mappings_config = self.config.source_mappings or []
 
-        for source in dremio_sources:
-            source_name = source.container_name
-            if isinstance(source.dremio_source_type, str):
-                source_type = source.dremio_source_type.lower()
-                root_path = source.root_path.lower() if source.root_path else ""
-                database_name = (
-                    source.database_name.lower() if source.database_name else ""
-                )
-                source_present = False
-                source_platform_name = source_name
-
-                for mapping in self.config.source_mappings or []:
-                    if re.search(mapping.source_name, source_type, re.IGNORECASE):
-                        source_platform_name = mapping.source_name.lower()
-
-                    datahub_source_type = (
-                        DremioToDataHubSourceTypeMapping.get_datahub_source_type(
-                            source_type
-                        )
-                    )
-
-                    if re.search(mapping.platform, datahub_source_type, re.IGNORECASE):
-                        source_platform_name = source_platform_name.lower()
-                        source_map[source_platform_name] = {
-                            "platform": mapping.platform,
-                            "source_name": mapping.source_name,
-                            "dremio_source_type": DremioToDataHubSourceTypeMapping.get_category(
-                                source_type,
-                            ),
-                            "root_path": root_path,
-                            "database_name": database_name,
-                            "platform_instance": mapping.platform_instance,
-                            "env": mapping.env,
-                        }
-                        source_present = True
-                        break
-
-                if not source_present:
-                    try:
-                        dremio_source_type = (
-                            DremioToDataHubSourceTypeMapping.get_category(source_type)
-                        )
-                    except Exception as exc:
-                        logger.info(
-                            f"Source {source_type} is not a standard Dremio source type. "
-                            f"Adding source_type {source_type} to mapping as database. Error: {exc}"
-                        )
-
-                        DremioToDataHubSourceTypeMapping.add_mapping(
-                            source_type, source_name
-                        )
-                        dremio_source_type = (
-                            DremioToDataHubSourceTypeMapping.get_category(source_type)
-                        )
-
-                    source_map[source_platform_name.lower()] = {
-                        "platform": source_type,
-                        "source_name": source_name,
-                        "dremio_source_type": dremio_source_type,
-                    }
-
-            else:
-                logger.error(
-                    f'Source "{source.container_name}" is broken. Containers will not be created for source.'
-                )
-                logger.error(
-                    f'No new cross-platform lineage will be emitted for source "{source.container_name}".'
-                )
-                logger.error("Fix this source in Dremio to fix this issue.")
+        source_map = build_dremio_source_map(dremio_sources, source_mappings_config)
+        logger.info(f"Full source map: {source_map}")
 
         return source_map
 
@@ -275,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()
+
+            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,
+                    )
 
-        # 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,
-                        )
+            # 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
@@ -385,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(
@@ -431,6 +384,7 @@ class DremioSource(StatefulIngestionSourceBase):
                     dremio_path=dataset_info.path,
                     dremio_dataset=dataset_info.resource_name,
                 )
+                logger.debug(f"Upstream dataset for {dataset_urn}: {upstream_urn}")
 
                 if upstream_urn:
                     upstream_lineage = UpstreamLineage(
@@ -467,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]
@@ -483,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
         ]
@@ -501,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:
@@ -548,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
@@ -596,25 +546,23 @@ class DremioSource(StatefulIngestionSourceBase):
         if not mapping:
             return None
 
-        platform = mapping.get("platform")
+        platform = mapping.platform
         if not platform:
             return None
 
-        platform_instance = mapping.get(
-            "platform_instance", self.config.platform_instance
-        )
-        env = mapping.get("env", self.config.env)
+        platform_instance = mapping.platform_instance
+        env = mapping.env or self.config.env
 
         root_path = ""
         database_name = ""
 
-        if mapping.get("dremio_source_type") == "file_object_storage":
-            if mapping.get("root_path"):
-                root_path = f"{mapping['root_path'][1:]}/"
+        if mapping.dremio_source_category == "file_object_storage":
+            if mapping.root_path:
+                root_path = f"{mapping.root_path[1:]}/"
             dremio_dataset = f"{root_path}{'/'.join(dremio_path[1:])}/{dremio_dataset}"
         else:
-            if mapping.get("database_name"):
-                database_name = f"{mapping['database_name']}."
+            if mapping.database_name:
+                database_name = f"{mapping.database_name}."
             dremio_dataset = (
                 f"{database_name}{'.'.join(dremio_path[1:])}.{dremio_dataset}"
             )
@@ -639,3 +587,68 @@ class DremioSource(StatefulIngestionSourceBase):
         Get the source report.
         """
         return self.report
+
+
+def build_dremio_source_map(
+    dremio_sources: Iterable[DremioSourceContainer],
+    source_mappings_config: List[DremioSourceMapping],
+) -> Dict[str, DremioSourceMapEntry]:
+    """
+    Builds a source mapping dictionary to support external lineage generation across
+    multiple Dremio sources, based on provided configuration mappings.
+
+    This method operates as follows:
+
+    Returns:
+        Dict[str, Dict]: A dictionary (`source_map`) where each key is a source name
+                        (lowercased) and each value is another entry containing:
+                        - `platform`: The source platform.
+                        - `source_name`: The source name.
+                        - `dremio_source_category`: The type mapped to DataHub,
+                        e.g., "database", "folder".
+                        - Optional `root_path`, `database_name`, `platform_instance`,
+                        and `env` if provided in the configuration.
+    Example:
+        This method is used internally within the class to generate mappings before
+        creating cross-platform lineage.
+
+    """
+    source_map = {}
+    for source in dremio_sources:
+        current_source_name = source.container_name
+
+        source_type = source.dremio_source_type.lower()
+        source_category = DremioToDataHubSourceTypeMapping.get_category(source_type)
+        datahub_platform = DremioToDataHubSourceTypeMapping.get_datahub_platform(
+            source_type
+        )
+        root_path = source.root_path.lower() if source.root_path else ""
+        database_name = source.database_name.lower() if source.database_name else ""
+        source_present = False
+
+        for mapping in source_mappings_config:
+            if mapping.source_name.lower() == current_source_name.lower():
+                source_map[current_source_name.lower()] = DremioSourceMapEntry(
+                    platform=mapping.platform,
+                    source_name=mapping.source_name,
+                    dremio_source_category=source_category,
+                    root_path=root_path,
+                    database_name=database_name,
+                    platform_instance=mapping.platform_instance,
+                    env=mapping.env,
+                )
+                source_present = True
+                break
+
+        if not source_present:
+            source_map[current_source_name.lower()] = DremioSourceMapEntry(
+                platform=datahub_platform,
+                source_name=current_source_name,
+                dremio_source_category=source_category,
+                root_path=root_path,
+                database_name=database_name,
+                platform_instance=None,
+                env=None,
+            )
+
+    return source_map