acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.2.0.1__py3-none-any.whl

datahub/ingestion/source/debug/datahub_debug.py

@@ -0,0 +1,300 @@
+import logging
+import socket
+import time
+from typing import Iterable, Optional
+from urllib.parse import urlparse
+
+import dns.exception
+import dns.resolver
+import requests
+
+from datahub.configuration.common import ConfigModel
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SupportStatus,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import Source, SourceReport
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+
+logger = logging.getLogger(__name__)
+
+
+class DataHubDebugSourceConfig(ConfigModel):
+    dns_probe_url: Optional[str] = None
+
+
+@platform_name("DataHubDebug")
+@config_class(DataHubDebugSourceConfig)
+@support_status(SupportStatus.TESTING)
+class DataHubDebugSource(Source):
+    """
+    DataHubDebugSource is helper to debug things in executor where ingestion is running.
+
+    This source can perform the following tasks:
+    1. Network probe of a URL. Different from test connection in sources as that is after source starts.
+
+    """
+
+    def __init__(self, ctx: PipelineContext, config: DataHubDebugSourceConfig):
+        self.ctx = ctx
+        self.config = config
+        self.report = SourceReport()
+        self.report.event_not_produced_warn = False
+
+    @classmethod
+    def create(cls, config_dict, ctx):
+        config = DataHubDebugSourceConfig.parse_obj(config_dict)
+        return cls(ctx, config)
+
+    def perform_dns_probe(self, url: str) -> None:
+        """
+        Perform comprehensive DNS probe and network connectivity tests.
+        Logs detailed information to help diagnose network issues.
+        """
+        logger.info(f"Starting DNS probe for URL: {url}")
+        logger.info("=" * 60)
+        logger.info(f"DNS PROBE REPORT FOR: {url}")
+        logger.info("=" * 60)
+
+        try:
+            # Parse the URL to extract hostname
+            parsed_url = urlparse(
+                url if url.startswith(("http://", "https://")) else f"http://{url}"
+            )
+            hostname = parsed_url.hostname or parsed_url.netloc
+            port = parsed_url.port or (443 if parsed_url.scheme == "https" else 80)
+
+            logger.info(f"Parsed hostname: {hostname}")
+            logger.info(f"Target port: {port}")
+            logger.info(f"URL scheme: {parsed_url.scheme}")
+            logger.info("-" * 60)
+
+            # Test 1: Enhanced DNS resolution with dnspython if available
+            logger.info("1. DNS RESOLUTION TEST")
+            self._dns_probe_with_dnspython(hostname)
+
+            logger.info("-" * 60)
+
+            # Test 2: HTTP/HTTPS connectivity test with requests if available
+            logger.info("2. HTTP CONNECTIVITY TEST")
+            self._http_probe_with_requests(url)
+
+            logger.info("-" * 60)
+
+            # Test 3: System network information
+            logger.info("3. SYSTEM NETWORK INFORMATION")
+            self._log_system_network_info()
+
+        except Exception as e:
+            logger.error(f"DNS probe failed with unexpected error: {e}", exc_info=True)
+
+        logger.info("=" * 60)
+        logger.info("DNS PROBE COMPLETED")
+        logger.info("=" * 60)
+
+    def _dns_probe_with_dnspython(self, hostname: str) -> None:
+        """Enhanced DNS probing using dnspython library"""
+        try:
+            # Test different record types
+            record_types = ["A", "AAAA", "CNAME", "MX"]
+
+            for record_type in record_types:
+                try:
+                    start_time = time.time()
+                    answers = dns.resolver.resolve(hostname, record_type)
+                    dns_time = time.time() - start_time
+
+                    logger.info(
+                        f"✓ {record_type} record resolution successful ({dns_time:.3f}s)"
+                    )
+                    for answer in answers:
+                        logger.info(f"  - {record_type}: {answer}")
+
+                except dns.resolver.NXDOMAIN:
+                    logger.info(f"✗ {record_type} record: Domain does not exist")
+                except dns.resolver.NoAnswer:
+                    logger.info(
+                        f"- {record_type} record: No answer (record type not available)"
+                    )
+                except dns.exception.Timeout:
+                    logger.error(f"✗ {record_type} record: DNS query timed out")
+                except Exception as e:
+                    logger.error(f"✗ {record_type} record query failed: {e}")
+
+            # Test different DNS servers
+            logger.info("Testing with different DNS servers:")
+            dns_servers = ["8.8.8.8", "1.1.1.1", "208.67.222.222"]
+
+            for dns_server in dns_servers:
+                try:
+                    resolver = dns.resolver.Resolver()
+                    resolver.nameservers = [dns_server]
+                    resolver.timeout = 5
+
+                    start_time = time.time()
+                    answers = resolver.resolve(hostname, "A")
+                    dns_time = time.time() - start_time
+
+                    logger.info(
+                        f"✓ DNS server {dns_server} responded ({dns_time:.3f}s)"
+                    )
+                    for answer in answers:
+                        logger.info(f"  - A: {answer}")
+
+                except Exception as e:
+                    logger.error(f"✗ DNS server {dns_server} failed: {e}")
+
+        except Exception as e:
+            logger.error(f"Enhanced DNS probe failed: {e}", exc_info=True)
+
+    def _http_probe_with_requests(self, url: str) -> None:
+        """HTTP connectivity test using requests library"""
+        try:
+            # Test with different timeouts and methods
+            timeout = 10
+            allow_redirects_head = True
+            allow_redirects_get = False
+
+            # Test HEAD request
+            try:
+                logger.info(f"Testing HEAD request with timeout {timeout}s")
+                start_time = time.time()
+
+                response = requests.head(
+                    url, timeout=timeout, allow_redirects=allow_redirects_head
+                )
+
+                request_time = time.time() - start_time
+
+                logger.info(f"✓ HEAD request successful ({request_time:.3f}s)")
+                logger.info(f"  Status code: {response.status_code}")
+                logger.info(
+                    f"  Response headers: {dict(list(response.headers.items())[:5])}"
+                )
+
+                if hasattr(response, "url") and response.url != url:
+                    logger.info(f"  Final URL after redirects: {response.url}")
+
+            except requests.exceptions.Timeout:
+                logger.error(f"✗ HEAD request timed out after {timeout}s")
+            except requests.exceptions.ConnectionError as e:
+                logger.error(f"✗ HEAD connection error: {e}")
+            except requests.exceptions.RequestException as e:
+                logger.error(f"✗ HEAD request failed: {e}")
+            except Exception as e:
+                logger.error(f"✗ HEAD unexpected error: {e}")
+
+            # Test GET request
+            try:
+                logger.info(f"Testing GET request with timeout {timeout}s")
+                start_time = time.time()
+
+                response = requests.get(
+                    url, timeout=timeout, allow_redirects=allow_redirects_get
+                )
+
+                request_time = time.time() - start_time
+
+                logger.info(f"✓ GET request successful ({request_time:.3f}s)")
+                logger.info(f"  Status code: {response.status_code}")
+                logger.info(
+                    f"  Response headers: {dict(list(response.headers.items())[:5])}"
+                )
+
+                if hasattr(response, "url") and response.url != url:
+                    logger.info(f"  Final URL after redirects: {response.url}")
+
+            except requests.exceptions.Timeout:
+                logger.error(f"✗ GET request timed out after {timeout}s")
+            except requests.exceptions.ConnectionError as e:
+                logger.error(f"✗ GET connection error: {e}")
+            except requests.exceptions.RequestException as e:
+                logger.error(f"✗ GET request failed: {e}")
+            except Exception as e:
+                logger.error(f"✗ GET unexpected error: {e}")
+
+        except Exception as e:
+            logger.error(f"HTTP probe failed: {e}", exc_info=True)
+
+    def _log_dns_troubleshooting(self) -> None:
+        """Log DNS troubleshooting information"""
+        logger.info("DNS TROUBLESHOOTING SUGGESTIONS:")
+        logger.info("- Check if the hostname is correct")
+        logger.info("- Verify DNS server configuration")
+        logger.info("- Check network connectivity")
+        logger.info("- Try using a different DNS server (8.8.8.8, 1.1.1.1)")
+        logger.info("- Check if there are firewall restrictions")
+
+    def _log_system_network_info(self) -> None:
+        """Log system network configuration information"""
+        try:
+            local_hostname = socket.gethostname()
+            logger.info(f"Local hostname: {local_hostname}")
+
+            try:
+                local_ips = socket.getaddrinfo(local_hostname, None)
+                logger.info("Local IP addresses:")
+                for addr_info in local_ips:
+                    if addr_info[0] in [socket.AF_INET, socket.AF_INET6]:
+                        family = "IPv4" if addr_info[0] == socket.AF_INET else "IPv6"
+                        logger.info(f"  - {addr_info[4][0]} ({family})")
+            except Exception as e:
+                logger.warning(f"Could not retrieve local IP addresses: {e}")
+
+            logger.info("DNS Server Connectivity:")
+            dns_servers = ["8.8.8.8", "1.1.1.1", "208.67.222.222"]
+            for dns_server in dns_servers:
+                try:
+                    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+                    sock.settimeout(5)
+                    result = sock.connect_ex((dns_server, 53))
+                    if result == 0:
+                        logger.info(f"  ✓ Can reach {dns_server}:53")
+                    else:
+                        logger.error(f"  ✗ Cannot reach {dns_server}:53")
+                    sock.close()
+                except Exception as e:
+                    logger.error(f"  ✗ Error testing {dns_server}:53 - {e}")
+
+        except Exception as e:
+            logger.warning(f"Could not gather system network info: {e}")
+
+    def _test_alternative_dns(self, hostname: str) -> None:
+        """Test hostname resolution using alternative methods"""
+        try:
+            families = [(socket.AF_INET, "IPv4"), (socket.AF_INET6, "IPv6")]
+
+            for family, family_name in families:
+                try:
+                    result = socket.getaddrinfo(hostname, None, family)
+                    if result:
+                        logger.info(f"✓ {family_name} resolution successful:")
+                        for addr_info in result[:3]:
+                            logger.info(f"  - {addr_info[4][0]}")
+                    else:
+                        logger.warning(
+                            f"✗ {family_name} resolution returned no results"
+                        )
+                except socket.gaierror:
+                    logger.error(f"✗ {family_name} resolution failed")
+                except Exception as e:
+                    logger.error(f"✗ {family_name} resolution error: {e}")
+
+        except Exception as e:
+            logger.error(f"Alternative DNS test failed: {e}")
+
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[MetadataWorkUnit]:
+        if self.config.dns_probe_url is not None:
+            # Perform DNS probe
+            logger.info(f"Performing DNS probe for: {self.config.dns_probe_url}")
+            self.perform_dns_probe(self.config.dns_probe_url)
+
+        yield from []
+
+    def get_report(self) -> SourceReport:
+        return self.report

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.