acryl-datahub 1.1.0.5rc8__py3-none-any.whl → 1.1.0.5rc10__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,6 +15,7 @@ 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,
@@ -23,6 +24,9 @@ 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__)
 
 
@@ -181,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)
@@ -192,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, (
@@ -215,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)
@@ -234,42 +240,45 @@ class DremioAPIOperations:
             "Credentials cannot be refreshed. Please check your username and password."
         )
 
-    def get(self, url: str) -> Dict:
-        """execute a get request on dremio"""
-        logger.debug(f"GET request to {self.base_url + url}")
-        self.report.api_calls_total += 1
-        self.report.api_calls_by_method_and_path["GET " + url] += 1
-
-        with PerfTimer() as timer:
-            response = self.session.get(
-                url=(self.base_url + url),
-                verify=self._verify,
-                timeout=self._timeout,
-            )
-            self.report.api_call_secs_by_method_and_path["GET " + url] += (
-                timer.elapsed_seconds()
-            )
-            # response.raise_for_status()  # Enabling this line, makes integration tests to fail
-            return response.json()
+    def _request(self, method: str, url: str, data: Union[str, None] = None) -> Dict:
+        """Send a request to the Dremio API."""
 
-    def post(self, url: str, data: str) -> Dict:
-        """execute a get request on dremio"""
-        logger.debug(f"POST request to {self.base_url + url}")
+        logger.debug(f"{method} request to {self.base_url + url}")
         self.report.api_calls_total += 1
-        self.report.api_calls_by_method_and_path["POST " + url] += 1
+        self.report.api_calls_by_method_and_path[f"{method} {url}"] += 1
 
         with PerfTimer() as timer:
-            response = self.session.post(
+            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["POST " + url] += (
+            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
-            return response.json()
+            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:
+        """Send a GET request to the Dremio API."""
+        return self._request("GET", url)
+
+    def post(self, url: str, data: str) -> Dict:
+        """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"""
@@ -489,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:

datahub/ingestion/source/mlflow.py

@@ -33,7 +33,10 @@ from datahub.ingestion.api.source import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_platforms import KNOWN_VALID_PLATFORM_NAMES
-from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
+from datahub.ingestion.source.common.subtypes import (
+    MLAssetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
@@ -138,6 +141,13 @@ class MLflowRegisteredModelStageInfo:
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for MLflow Registered Models and Model Versions",
 )
+@capability(
+    SourceCapability.CONTAINERS,
+    "Extract ML experiments",
+    subtype_modifier=[
+        SourceCapabilityModifier.MLFLOW_EXPERIMENT,
+    ],
+)
 @capability(SourceCapability.TAGS, "Extract tags for MLflow Registered Model Stages")
 class MLflowSource(StatefulIngestionSourceBase):
     platform = "mlflow"

datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -119,6 +119,20 @@ class SnowflakeQueriesExtractorConfig(ConfigModel):
     include_query_usage_statistics: bool = True
     include_operations: bool = True
 
+    push_down_database_pattern_access_history: bool = pydantic.Field(
+        default=False,
+        description="If enabled, pushes down database pattern filtering to the access_history table for improved performance. "
+        "This filters on the accessed objects in access_history.",
+    )
+
+    additional_database_names_allowlist: List[str] = pydantic.Field(
+        default=[],
+        description="Additional database names (no pattern matching) to be included in the access_history filter. "
+        "Only applies if push_down_database_pattern_access_history=True. "
+        "These databases will be included in the filter being pushed down regardless of database_pattern settings."
+        "This may be required in the case of _eg_ temporary tables being created in a different database than the ones in the database_name patterns.",
+    )
+
     query_dedup_strategy: QueryDedupStrategyType = QueryDedupStrategyType.STANDARD
 
 
@@ -383,6 +397,12 @@ class SnowflakeQueriesExtractor(SnowflakeStructuredReportMixin, Closeable):
             bucket_duration=self.config.window.bucket_duration,
             deny_usernames=self.config.pushdown_deny_usernames,
             dedup_strategy=self.config.query_dedup_strategy,
+            database_pattern=self.filters.filter_config.database_pattern
+            if self.config.push_down_database_pattern_access_history
+            else None,
+            additional_database_names=self.config.additional_database_names_allowlist
+            if self.config.push_down_database_pattern_access_history
+            else None,
         ).build_enriched_query_log_query()
 
         with self.structured_reporter.report_exc(
@@ -723,6 +743,8 @@ class QueryLogQueryBuilder:
         deny_usernames: Optional[List[str]],
         max_tables_per_query: int = 20,
         dedup_strategy: QueryDedupStrategyType = QueryDedupStrategyType.STANDARD,
+        database_pattern: Optional[AllowDenyPattern] = None,
+        additional_database_names: Optional[List[str]] = None,
     ):
         self.start_time = start_time
         self.end_time = end_time
@@ -736,9 +758,113 @@ class QueryLogQueryBuilder:
             user_not_in = ",".join(f"'{user.upper()}'" for user in deny_usernames)
             self.users_filter = f"user_name NOT IN ({user_not_in})"
 
+        self.access_history_database_filter = (
+            self._build_access_history_database_filter_condition(
+                database_pattern, additional_database_names
+            )
+        )
+
         self.time_bucket_size = bucket_duration.value
         assert self.time_bucket_size in ("HOUR", "DAY", "MONTH")
 
+    def _build_access_history_database_filter_condition(
+        self,
+        database_pattern: Optional[AllowDenyPattern],
+        additional_database_names: Optional[List[str]] = None,
+    ) -> str:
+        """
+        Build a SQL WHERE condition for database filtering in access_history based on AllowDenyPattern.
+
+        IMPORTANT: This function handles the fundamental difference between DML and DDL operations in Snowflake's
+        access_history table:
+
+        - DML Operations (SELECT, INSERT, UPDATE, DELETE, etc.): Store accessed/modified objects in the
+          `direct_objects_accessed` and `objects_modified` arrays
+        - DDL Operations (CREATE, ALTER, DROP, RENAME, etc.): Store modified objects in the
+          `object_modified_by_ddl` field (single object, not an array)
+
+        Without checking `object_modified_by_ddl`, DDL operations like "ALTER TABLE person_info RENAME TO person_info_final"
+        would be incorrectly filtered out because they don't populate the DML arrays, causing missing lineage
+        and operational metadata.
+
+        Args:
+            database_pattern: The AllowDenyPattern configuration for database filtering
+            additional_database_names: Additional database names to always include (no pattern matching)
+
+        Returns:
+            A SQL WHERE condition string, or "TRUE" if no filtering should be applied
+        """
+        if not database_pattern and not additional_database_names:
+            return "TRUE"
+
+        # Build the database filter conditions for pattern matching
+        # Note: Using UPPER() + RLIKE for case-insensitive matching is more performant than REGEXP_LIKE with 'i' flag
+        database_filter_parts = []
+
+        if database_pattern:
+            allow_patterns = database_pattern.allow
+            deny_patterns = database_pattern.deny
+
+            # Add allow patterns (if not the default "allow all")
+            if allow_patterns and allow_patterns != [".*"]:
+                allow_conditions = []
+                for pattern in allow_patterns:
+                    # Escape single quotes that might be present in the regex pattern
+                    escaped_pattern = pattern.replace("'", "''")
+                    allow_conditions.append(
+                        f"SPLIT_PART(UPPER(o:objectName), '.', 1) RLIKE '{escaped_pattern}'"
+                    )
+                if allow_conditions:
+                    database_filter_parts.append(f"({' OR '.join(allow_conditions)})")
+
+            # Add deny patterns
+            if deny_patterns:
+                deny_conditions = []
+                for pattern in deny_patterns:
+                    # Escape single quotes that might be present in the regex pattern
+                    escaped_pattern = pattern.replace("'", "''")
+                    deny_conditions.append(
+                        f"SPLIT_PART(UPPER(o:objectName), '.', 1) NOT RLIKE '{escaped_pattern}'"
+                    )
+                if deny_conditions:
+                    database_filter_parts.append(f"({' AND '.join(deny_conditions)})")
+
+        # Add additional database names (exact matches)
+        if additional_database_names:
+            additional_db_conditions = []
+            for db_name in additional_database_names:
+                # Escape single quotes
+                escaped_db_name = db_name.replace("'", "''")
+                additional_db_conditions.append(
+                    f"SPLIT_PART(UPPER(o:objectName), '.', 1) = '{escaped_db_name.upper()}'"
+                )
+            if additional_db_conditions:
+                database_filter_parts.append(
+                    f"({' OR '.join(additional_db_conditions)})"
+                )
+
+        if database_filter_parts:
+            database_filter_condition = " AND ".join(database_filter_parts)
+
+            # Build a condition that checks if any objects in the arrays match the database pattern
+            # This implements "at least one" matching behavior: queries are allowed if they touch
+            # at least one database that matches the pattern, even if they also touch other databases
+            # Use ARRAY_SIZE with FILTER which is more compatible with Snowflake
+            direct_objects_condition = f"ARRAY_SIZE(FILTER(direct_objects_accessed, o -> {database_filter_condition})) > 0"
+            objects_modified_condition = f"ARRAY_SIZE(FILTER(objects_modified, o -> {database_filter_condition})) > 0"
+
+            # CRITICAL: Handle DDL operations by checking object_modified_by_ddl field
+            # DDL operations like ALTER TABLE RENAME store their data here instead of in the arrays
+            # We need to adapt the filter condition for a single object rather than an array
+            ddl_filter_condition = database_filter_condition.replace(
+                "o:objectName", "object_modified_by_ddl:objectName"
+            )
+            object_modified_by_ddl_condition = f"({ddl_filter_condition})"
+
+            return f"({direct_objects_condition} OR {objects_modified_condition} OR {object_modified_by_ddl_condition})"
+        else:
+            return "TRUE"
+
     def _query_fingerprinted_queries(self):
         if self.dedup_strategy == QueryDedupStrategyType.STANDARD:
             secondary_fingerprint_sql = """
@@ -828,6 +954,7 @@ fingerprinted_queries as (
         AND query_id IN (
             SELECT query_id FROM deduplicated_queries
         )
+        AND {self.access_history_database_filter}
 )
 , filtered_access_history AS (
     -- TODO: Add table filter clause.

datahub/ingestion/source/sql/sql_common.py

@@ -593,6 +593,10 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
                 )
 
         # Generate workunit for aggregated SQL parsing results
+        yield from self._generate_aggregator_workunits()
+
+    def _generate_aggregator_workunits(self) -> Iterable[MetadataWorkUnit]:
+        """Generate work units from SQL parsing aggregator. Can be overridden by subclasses."""
         for mcp in self.aggregator.gen_metadata():
             yield mcp.as_workunit()