acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.0.0.1rc2__py3-none-any.whl

datahub/ingestion/source/data_lake_common/path_spec.py

@@ -11,7 +11,7 @@ from cached_property import cached_property
 from pydantic.fields import Field
 from wcmatch import pathlib
 
-from datahub.configuration.common import ConfigModel
+from datahub.configuration.common import AllowDenyPattern, ConfigModel
 from datahub.ingestion.source.aws.s3_util import is_s3_uri
 from datahub.ingestion.source.azure.abs_utils import is_abs_uri
 from datahub.ingestion.source.gcs.gcs_utils import is_gcs_uri
@@ -145,6 +145,11 @@ class PathSpec(ConfigModel):
         description="Include hidden folders in the traversal (folders starting with . or _",
     )
 
+    tables_filter_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="The tables_filter_pattern configuration field uses regular expressions to filter the tables part of the Pathspec for ingestion, allowing fine-grained control over which tables are included or excluded based on specified patterns. The default setting allows all tables.",
+    )
+
     def is_path_hidden(self, path: str) -> bool:
         # Split the path into directories and filename
         dirs, filename = os.path.split(path)
@@ -177,6 +182,12 @@ class PathSpec(ConfigModel):
                 ):
                     return False
         logger.debug(f"{path} is not excluded")
+
+        table_name, _ = self.extract_table_name_and_path(path)
+        if not self.tables_filter_pattern.allowed(table_name):
+            return False
+        logger.debug(f"{path} is passed table name check")
+
         ext = os.path.splitext(path)[1].strip(".")
 
         if not ignore_ext:
@@ -218,6 +229,15 @@ class PathSpec(ConfigModel):
                     exclude_path.rstrip("/"), flags=pathlib.GLOBSTAR
                 ):
                     return False
+
+        file_name_pattern = self.include.rsplit("/", 1)[1]
+        table_name, _ = self.extract_table_name_and_path(
+            os.path.join(path, file_name_pattern)
+        )
+        if not self.tables_filter_pattern.allowed(table_name):
+            return False
+        logger.debug(f"{path} is passed table name check")
+
         return True
 
     @classmethod

datahub/ingestion/source/dbt/dbt_common.py

@@ -4,7 +4,7 @@ from abc import abstractmethod
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import auto
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
 
 import more_itertools
 import pydantic
@@ -849,7 +849,7 @@ class DBTSourceBase(StatefulIngestionSourceBase):
         test_nodes: List[DBTNode],
         extra_custom_props: Dict[str, str],
         all_nodes_map: Dict[str, DBTNode],
-    ) -> Iterable[MetadataWorkUnit]:
+    ) -> Iterable[MetadataChangeProposalWrapper]:
         for node in sorted(test_nodes, key=lambda n: n.dbt_name):
             upstreams = get_upstreams_for_test(
                 test_node=node,
@@ -902,7 +902,7 @@ class DBTSourceBase(StatefulIngestionSourceBase):
                     yield MetadataChangeProposalWrapper(
                         entityUrn=assertion_urn,
                         aspect=self._make_data_platform_instance_aspect(),
-                    ).as_workunit()
+                    )
 
                     yield make_assertion_from_test(
                         custom_props,
@@ -949,7 +949,9 @@ class DBTSourceBase(StatefulIngestionSourceBase):
             ),
         )
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         if self.config.write_semantics == "PATCH":
             self.ctx.require_graph("Using dbt with write_semantics=PATCH")
 

datahub/ingestion/source/dbt/dbt_core.py

@@ -343,6 +343,9 @@ class DBTRunResult(BaseModel):
     def timing_map(self) -> Dict[str, DBTRunTiming]:
         return {x.name: x for x in self.timing if x.name}
 
+    def has_success_status(self) -> bool:
+        return self.status in ("pass", "success")
+
 
 class DBTRunMetadata(BaseModel):
     dbt_schema_version: str
@@ -355,12 +358,7 @@ def _parse_test_result(
     dbt_metadata: DBTRunMetadata,
     run_result: DBTRunResult,
 ) -> Optional[DBTTestResult]:
-    if run_result.status == "success":
-        # This was probably a docs generate run result, so this isn't actually
-        # a test result.
-        return None
-
-    if run_result.status != "pass":
+    if not run_result.has_success_status():
         native_results = {"message": run_result.message or ""}
         if run_result.failures:
             native_results.update({"failures": str(run_result.failures)})

datahub/ingestion/source/dbt/dbt_tests.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union
 
 from datahub.emitter import mce_builder
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.schema_classes import (
     AssertionInfoClass,
     AssertionResultClass,
@@ -43,6 +42,9 @@ class DBTTestResult:
 
     native_results: Dict[str, str]
 
+    def has_success_status(self) -> bool:
+        return self.status in ("pass", "success")
+
 
 def _get_name_for_relationship_test(kw_args: Dict[str, str]) -> Optional[str]:
     """
@@ -157,7 +159,7 @@ def make_assertion_from_test(
     node: "DBTNode",
     assertion_urn: str,
     upstream_urn: str,
-) -> MetadataWorkUnit:
+) -> MetadataChangeProposalWrapper:
     assert node.test_info
     qualified_test_name = node.test_info.qualified_test_name
     column_name = node.test_info.column_name
@@ -231,7 +233,7 @@ def make_assertion_from_test(
     return MetadataChangeProposalWrapper(
         entityUrn=assertion_urn,
         aspect=assertion_info,
-    ).as_workunit()
+    )
 
 
 def make_assertion_result_from_test(
@@ -240,7 +242,7 @@ def make_assertion_result_from_test(
     assertion_urn: str,
     upstream_urn: str,
     test_warnings_are_errors: bool,
-) -> MetadataWorkUnit:
+) -> MetadataChangeProposalWrapper:
     assertionResult = AssertionRunEventClass(
         timestampMillis=int(test_result.execution_time.timestamp() * 1000.0),
         assertionUrn=assertion_urn,
@@ -249,7 +251,7 @@ def make_assertion_result_from_test(
         result=AssertionResultClass(
             type=(
                 AssertionResultTypeClass.SUCCESS
-                if test_result.status == "pass"
+                if test_result.has_success_status()
                 or (not test_warnings_are_errors and test_result.status == "warn")
                 else AssertionResultTypeClass.FAILURE
             ),
@@ -261,4 +263,4 @@ def make_assertion_result_from_test(
     return MetadataChangeProposalWrapper(
         entityUrn=assertion_urn,
         aspect=assertionResult,
-    ).as_workunit()
+    )

datahub/ingestion/source/dremio/dremio_datahub_source_mapping.py

@@ -66,7 +66,7 @@ class DremioToDataHubSourceTypeMapping:
     }
 
     @staticmethod
-    def get_datahub_source_type(dremio_source_type: str) -> str:
+    def get_datahub_platform(dremio_source_type: str) -> str:
         """
         Return the DataHub source type.
         """

datahub/ingestion/source/dremio/dremio_entities.py

@@ -294,7 +294,7 @@ class DremioContainer:
         )
 
 
-class DremioSource(DremioContainer):
+class DremioSourceContainer(DremioContainer):
     subclass: str = "Dremio Source"
     dremio_source_type: str
     root_path: Optional[str]
@@ -337,7 +337,7 @@ class DremioCatalog:
         self.dremio_api = dremio_api
         self.edition = dremio_api.edition
         self.datasets: Deque[DremioDataset] = deque()
-        self.sources: Deque[DremioSource] = deque()
+        self.sources: Deque[DremioSourceContainer] = deque()
         self.spaces: Deque[DremioSpace] = deque()
         self.folders: Deque[DremioFolder] = deque()
         self.glossary_terms: Deque[DremioGlossaryTerm] = deque()
@@ -380,12 +380,13 @@ class DremioCatalog:
                 container_type = container.get("container_type")
                 if container_type == DremioEntityContainerType.SOURCE:
                     self.sources.append(
-                        DremioSource(
+                        DremioSourceContainer(
                             container_name=container.get("name"),
                             location_id=container.get("id"),
                             path=[],
                             api_operations=self.dremio_api,
-                            dremio_source_type=container.get("source_type"),
+                            dremio_source_type=container.get("source_type")
+                            or "unknown",
                             root_path=container.get("root_path"),
                             database_name=container.get("database_name"),
                         )
@@ -426,7 +427,7 @@ class DremioCatalog:
         self.set_containers()
         return deque(itertools.chain(self.sources, self.spaces, self.folders))
 
-    def get_sources(self) -> Deque[DremioSource]:
+    def get_sources(self) -> Deque[DremioSourceContainer]:
         self.set_containers()
         return self.sources
 

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 (
@@ -28,7 +27,10 @@ from datahub.ingestion.source.dremio.dremio_api import (
     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 +41,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
@@ -65,6 +68,17 @@ 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)
@@ -112,7 +126,7 @@ class DremioSource(StatefulIngestionSourceBase):
         self.default_db = "dremio"
         self.config = config
         self.report = DremioSourceReport()
-        self.source_map: Dict[str, Dict] = defaultdict()
+        self.source_map: Dict[str, DremioSourceMapEntry] = dict()
 
         # Initialize API operations
         dremio_api = DremioAPIOperations(self.config, self.report)
@@ -152,111 +166,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
 
@@ -431,6 +346,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(
@@ -596,25 +512,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 +553,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