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

datahub/ingestion/source/delta_lake/source.py

@@ -29,6 +29,7 @@ from datahub.ingestion.source.aws.s3_util import (
     get_key_prefix,
     strip_s3_prefix,
 )
+from datahub.ingestion.source.common.subtypes import SourceCapabilityModifier
 from datahub.ingestion.source.data_lake_common.data_lake_utils import ContainerWUCreator
 from datahub.ingestion.source.delta_lake.config import DeltaLakeSourceConfig
 from datahub.ingestion.source.delta_lake.delta_lake_utils import (
@@ -85,6 +86,13 @@ OPERATION_STATEMENT_TYPES = {
 @config_class(DeltaLakeSourceConfig)
 @support_status(SupportStatus.INCUBATING)
 @capability(SourceCapability.TAGS, "Can extract S3 object/bucket tags if enabled")
+@capability(
+    SourceCapability.CONTAINERS,
+    "Enabled by default",
+    subtype_modifier=[
+        SourceCapabilityModifier.FOLDER,
+    ],
+)
 class DeltaLakeSource(StatefulIngestionSourceBase):
     """
     This plugin extracts:

datahub/ingestion/source/dremio/dremio_api.py

@@ -7,7 +7,7 @@ from collections import defaultdict
 from enum import Enum
 from itertools import product
 from time import sleep, time
-from typing import Any, Deque, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Union
 from urllib.parse import quote
 
 import requests
@@ -15,12 +15,17 @@ from requests.adapters import HTTPAdapter
 from urllib3 import Retry
 from urllib3.exceptions import InsecureRequestWarning
 
+from datahub.emitter.request_helper import make_curl_command
 from datahub.ingestion.source.dremio.dremio_config import DremioSourceConfig
 from datahub.ingestion.source.dremio.dremio_datahub_source_mapping import (
     DremioToDataHubSourceTypeMapping,
 )
 from datahub.ingestion.source.dremio.dremio_reporting import DremioSourceReport
 from datahub.ingestion.source.dremio.dremio_sql_queries import DremioSQLQueries
+from datahub.utilities.perf_timer import PerfTimer
+
+if TYPE_CHECKING:
+    from datahub.ingestion.source.dremio.dremio_entities import DremioContainer
 
 logger = logging.getLogger(__name__)
 
@@ -54,6 +59,8 @@ class DremioAPIOperations:
         self.deny_schema_pattern: List[str] = connection_args.schema_pattern.deny
         self._max_workers: int = connection_args.max_workers
         self.is_dremio_cloud = connection_args.is_dremio_cloud
+        self.start_time = connection_args.start_time
+        self.end_time = connection_args.end_time
         self.report = report
         self.session = requests.Session()
         if connection_args.is_dremio_cloud:
@@ -178,6 +185,7 @@ class DremioAPIOperations:
             self.session.headers.update(
                 {"Authorization": f"Bearer {connection_args.password}"}
             )
+            logger.debug("Configured Dremio cloud API session to use PAT")
             return
 
         # On-prem Dremio authentication (PAT or Basic Auth)
@@ -189,6 +197,7 @@ class DremioAPIOperations:
                             "Authorization": f"Bearer {connection_args.password}",
                         }
                     )
+                    logger.debug("Configured Dremio API session to use PAT")
                     return
                 else:
                     assert connection_args.username and connection_args.password, (
@@ -212,10 +221,10 @@ class DremioAPIOperations:
                     response.raise_for_status()
                     token = response.json().get("token")
                     if token:
+                        logger.debug("Exchanged username and password for Dremio token")
                         self.session.headers.update(
                             {"Authorization": f"_dremio{token}"}
                         )
-
                         return
                     else:
                         self.report.failure("Failed to authenticate", login_url)
@@ -231,49 +240,76 @@ class DremioAPIOperations:
             "Credentials cannot be refreshed. Please check your username and password."
         )
 
+    def _request(self, method: str, url: str, data: Union[str, None] = None) -> Dict:
+        """Send a request to the Dremio API."""
+
+        logger.debug(f"{method} request to {self.base_url + url}")
+        self.report.api_calls_total += 1
+        self.report.api_calls_by_method_and_path[f"{method} {url}"] += 1
+
+        with PerfTimer() as timer:
+            response = self.session.request(
+                method=method,
+                url=(self.base_url + url),
+                data=data,
+                verify=self._verify,
+                timeout=self._timeout,
+            )
+            self.report.api_call_secs_by_method_and_path[f"{method} {url}"] += (
+                timer.elapsed_seconds()
+            )
+            # response.raise_for_status()  # Enabling this line, makes integration tests to fail
+            try:
+                return response.json()
+            except requests.exceptions.JSONDecodeError as e:
+                logger.info(
+                    f"On {method} request to {url}, failed to parse JSON from response (status {response.status_code}): {response.text}"
+                )
+                logger.debug(
+                    f"Request curl equivalent: {make_curl_command(self.session, method, url, data)}"
+                )
+                raise DremioAPIException(
+                    f"Failed to parse JSON from response (status {response.status_code}): {response.text}"
+                ) from e
+
     def get(self, url: str) -> Dict:
-        """execute a get request on dremio"""
-        response = self.session.get(
-            url=(self.base_url + url),
-            verify=self._verify,
-            timeout=self._timeout,
-        )
-        return response.json()
+        """Send a GET request to the Dremio API."""
+        return self._request("GET", url)
 
     def post(self, url: str, data: str) -> Dict:
-        """execute a get request on dremio"""
-        response = self.session.post(
-            url=(self.base_url + url),
-            data=data,
-            verify=self._verify,
-            timeout=self._timeout,
-        )
-        return response.json()
+        """Send a POST request to the Dremio API."""
+        return self._request("POST", url, data=data)
 
     def execute_query(self, query: str, timeout: int = 3600) -> List[Dict[str, Any]]:
         """Execute SQL query with timeout and error handling"""
         try:
-            response = self.post(url="/sql", data=json.dumps({"sql": query}))
+            with PerfTimer() as timer:
+                logger.info(f"Executing query: {query}")
+                response = self.post(url="/sql", data=json.dumps({"sql": query}))
 
-            if "errorMessage" in response:
-                self.report.failure(
-                    message="SQL Error", context=f"{response['errorMessage']}"
-                )
-                raise DremioAPIException(f"SQL Error: {response['errorMessage']}")
+                if "errorMessage" in response:
+                    self.report.failure(
+                        message="SQL Error", context=f"{response['errorMessage']}"
+                    )
+                    raise DremioAPIException(f"SQL Error: {response['errorMessage']}")
 
-            job_id = response["id"]
+                job_id = response["id"]
 
-            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
-                future = executor.submit(self.fetch_results, job_id)
-                try:
-                    return future.result(timeout=timeout)
-                except concurrent.futures.TimeoutError:
-                    self.cancel_query(job_id)
-                    raise DremioAPIException(
-                        f"Query execution timed out after {timeout} seconds"
-                    ) from None
-                except RuntimeError as e:
-                    raise DremioAPIException() from e
+                with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
+                    future = executor.submit(self.fetch_results, job_id)
+                    try:
+                        result = future.result(timeout=timeout)
+                        logger.info(
+                            f"Query executed in {timer.elapsed_seconds()} seconds with {len(result)} results"
+                        )
+                        return result
+                    except concurrent.futures.TimeoutError:
+                        self.cancel_query(job_id)
+                        raise DremioAPIException(
+                            f"Query execution timed out after {timeout} seconds"
+                        ) from None
+                    except RuntimeError as e:
+                        raise DremioAPIException() from e
 
         except requests.RequestException as e:
             raise DremioAPIException("Error executing query") from e
@@ -462,7 +498,9 @@ class DremioAPIOperations:
         pattern_str = "|".join(f"({p})" for p in patterns)
         return f"AND {operator}({field}, '{pattern_str}')"
 
-    def get_all_tables_and_columns(self, containers: Deque) -> List[Dict]:
+    def get_all_tables_and_columns(
+        self, containers: Deque["DremioContainer"]
+    ) -> List[Dict]:
         if self.edition == DremioEdition.ENTERPRISE:
             query_template = DremioSQLQueries.QUERY_DATASETS_EE
         elif self.edition == DremioEdition.CLOUD:
@@ -603,10 +641,25 @@ class DremioAPIOperations:
         return parents_list
 
     def extract_all_queries(self) -> List[Dict[str, Any]]:
+        # Convert datetime objects to string format for SQL queries
+        start_timestamp_str = None
+        end_timestamp_str = None
+
+        if self.start_time:
+            start_timestamp_str = self.start_time.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+        if self.end_time:
+            end_timestamp_str = self.end_time.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
+
         if self.edition == DremioEdition.CLOUD:
-            jobs_query = DremioSQLQueries.QUERY_ALL_JOBS_CLOUD
+            jobs_query = DremioSQLQueries.get_query_all_jobs_cloud(
+                start_timestamp_millis=start_timestamp_str,
+                end_timestamp_millis=end_timestamp_str,
+            )
         else:
-            jobs_query = DremioSQLQueries.QUERY_ALL_JOBS
+            jobs_query = DremioSQLQueries.get_query_all_jobs(
+                start_timestamp_millis=start_timestamp_str,
+                end_timestamp_millis=end_timestamp_str,
+            )
 
         return self.execute_query(query=jobs_query)
 
@@ -685,6 +738,27 @@ class DremioAPIOperations:
 
         return any(re.match(regex_pattern, path, re.IGNORECASE) for path in paths)
 
+    def _could_match_pattern(self, pattern: str, path_components: List[str]) -> bool:
+        """
+        Check if a container path could potentially match a schema pattern.
+        This handles hierarchical path matching for container filtering.
+        """
+        if pattern == ".*":
+            return True
+
+        current_path = ".".join(path_components)
+
+        # Handle simple .* patterns (like "a.b.c.*")
+        if pattern.endswith(".*") and not any(c in pattern for c in "^$[](){}+?\\"):
+            # Simple dotstar pattern - check prefix matching
+            pattern_prefix = pattern[:-2]  # Remove ".*"
+            return current_path.lower().startswith(
+                pattern_prefix.lower()
+            ) or pattern_prefix.lower().startswith(current_path.lower())
+        else:
+            # Complex regex pattern - use existing regex matching logic
+            return self._check_pattern_match(pattern, [current_path], allow_prefix=True)
+
     def should_include_container(self, path: List[str], name: str) -> bool:
         """
         Helper method to check if a container should be included based on schema patterns.
@@ -711,41 +785,8 @@ class DremioAPIOperations:
 
         # Check allow patterns
         for pattern in self.allow_schema_pattern:
-            # For patterns with wildcards, check if this path is a parent of the pattern
-            if "*" in pattern:
-                pattern_parts = pattern.split(".")
-                path_parts = path_components
-
-                # If pattern has exact same number of parts, check each component
-                if len(pattern_parts) == len(path_parts):
-                    matches = True
-                    for p_part, c_part in zip(pattern_parts, path_parts):
-                        if p_part != "*" and p_part.lower() != c_part.lower():
-                            matches = False
-                            break
-                    if matches:
-                        self.report.report_container_scanned(full_path)
-                        return True
-                # Otherwise check if current path is prefix match
-                else:
-                    # Remove the trailing wildcard if present
-                    if pattern_parts[-1] == "*":
-                        pattern_parts = pattern_parts[:-1]
-
-                    for i in range(len(path_parts)):
-                        current_path = ".".join(path_parts[: i + 1])
-                        pattern_prefix = ".".join(pattern_parts[: i + 1])
-
-                        if pattern_prefix.startswith(current_path):
-                            self.report.report_container_scanned(full_path)
-                            return True
-
-            # Direct pattern matching
-            if self._check_pattern_match(
-                pattern=pattern,
-                paths=[full_path],
-                allow_prefix=True,
-            ):
+            # Check if current path could potentially match this pattern
+            if self._could_match_pattern(pattern, path_components):
                 self.report.report_container_scanned(full_path)
                 return True
 

datahub/ingestion/source/dremio/dremio_aspects.py

@@ -14,6 +14,7 @@ from datahub.emitter.mce_builder import (
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import ContainerKey
 from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.common.subtypes import DatasetContainerSubTypes
 from datahub.ingestion.source.dremio.dremio_entities import (
     DremioContainer,
     DremioDataset,
@@ -364,9 +365,9 @@ class DremioAspects:
     ) -> Optional[BrowsePathsV2Class]:
         paths = []
 
-        if entity.subclass == "Dremio Space":
+        if entity.subclass == DatasetContainerSubTypes.DREMIO_SPACE.value:
             paths.append(BrowsePathEntryClass(id="Spaces"))
-        elif entity.subclass == "Dremio Source":
+        elif entity.subclass == DatasetContainerSubTypes.DREMIO_SOURCE.value:
             paths.append(BrowsePathEntryClass(id="Sources"))
         if paths:
             return BrowsePathsV2Class(path=paths)

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_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_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.