acryl-datahub 1.1.0.4rc2__py3-none-any.whl → 1.1.0.5__py3-none-any.whl

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_config.py

@@ -9,6 +9,7 @@ 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,
@@ -118,6 +119,7 @@ class DremioSourceMapping(EnvConfigMixin, PlatformInstanceConfigMixin, ConfigMod
 class DremioSourceConfig(
     DremioConnectionConfig,
     StatefulIngestionConfigBase,
+    BaseTimeWindowConfig,
     EnvConfigMixin,
     PlatformInstanceConfigMixin,
 ):

datahub/ingestion/source/dremio/dremio_reporting.py

@@ -1,22 +1,43 @@
-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,
+    IngestionStageReport,
+    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

@@ -51,13 +51,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,
+    PROFILING,
+)
 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,
@@ -89,6 +93,7 @@ class DremioSourceMapEntry:
 @capability(SourceCapability.LINEAGE_COARSE, "Enabled by default")
 @capability(SourceCapability.OWNERSHIP, "Enabled by default")
 @capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
+@capability(SourceCapability.USAGE_STATS, "Enabled by default to get usage stats")
 class DremioSource(StatefulIngestionSourceBase):
     """
     This plugin integrates with Dremio to extract and ingest metadata into DataHub.
@@ -126,6 +131,13 @@ class DremioSource(StatefulIngestionSourceBase):
         self.default_db = "dremio"
         self.config = config
         self.report = DremioSourceReport()
+
+        # Set time window for query lineage extraction
+        self.report.window_start_time, self.report.window_end_time = (
+            self.config.start_time,
+            self.config.end_time,
+        )
+
         self.source_map: Dict[str, DremioSourceMapEntry] = dict()
 
         # Initialize API operations
@@ -154,6 +166,7 @@ class DremioSource(StatefulIngestionSourceBase):
             generate_operations=True,
             usage_config=self.config.usage,
         )
+        self.report.sql_aggregator = self.sql_parsing_aggregator.report
 
         # For profiling
         self.profiler = DremioProfiler(config, self.report, dremio_api)
@@ -190,84 +203,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_stage(PROFILING),
+                    ThreadPoolExecutor(
+                        max_workers=self.config.profiling.max_workers
+                    ) as executor,
+                ):
+                    future_to_dataset = {
+                        executor.submit(self.generate_profiles, dataset): dataset
+                        for dataset in datasets
+                    }
+
+                    for future in as_completed(future_to_dataset):
+                        dataset_info = future_to_dataset[future]
+                        try:
+                            yield from future.result()
+                        except Exception as exc:
+                            self.report.profiling_skipped_other[
+                                dataset_info.resource_name
+                            ] += 1
+                            self.report.report_failure(
+                                message="Failed to profile dataset",
+                                context=f"{'.'.join(dataset_info.path)}.{dataset_info.resource_name}",
+                                exc=exc,
+                            )
 
     def process_container(
         self, container_info: DremioContainer
@@ -388,8 +405,7 @@ class DremioSource(StatefulIngestionSourceBase):
             env=self.config.env,
             platform_instance=self.config.platform_instance,
         )
-        with self.report.new_stage(f"{dataset_info.resource_name}: {PROFILING}"):
-            yield from self.profiler.get_workunits(dataset_info, dataset_urn)
+        yield from self.profiler.get_workunits(dataset_info, dataset_urn)
 
     def generate_view_lineage(
         self, dataset_urn: str, parents: List[str]
@@ -417,11 +433,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: