acryl-datahub 1.2.0.7rc3__py3-none-any.whl → 1.2.0.8rc1__py3-none-any.whl

datahub/ingestion/source/sql/athena.py

@@ -73,6 +73,11 @@ except ImportError:
 
 logger = logging.getLogger(__name__)
 
+# Precompiled regex for SQL identifier validation
+# Athena identifiers can only contain lowercase letters, numbers, underscore, and period (for complex types)
+# Note: Athena automatically converts uppercase to lowercase, but we're being strict for security
+_IDENTIFIER_PATTERN = re.compile(r"^[a-zA-Z0-9_.]+$")
+
 assert STRUCT, "required type modules are not available"
 register_custom_type(STRUCT, RecordTypeClass)
 register_custom_type(MapType, MapTypeClass)
@@ -510,20 +515,76 @@ class AthenaSource(SQLAlchemySource):
             return [schema for schema in schemas if schema == athena_config.database]
         return schemas
 
+    @classmethod
+    def _sanitize_identifier(cls, identifier: str) -> str:
+        """Sanitize SQL identifiers to prevent injection attacks.
+
+        Args:
+            identifier: The SQL identifier to sanitize
+
+        Returns:
+            Sanitized identifier safe for SQL queries
+
+        Raises:
+            ValueError: If identifier contains unsafe characters
+        """
+        if not identifier:
+            raise ValueError("Identifier cannot be empty")
+
+        # Allow only alphanumeric characters, underscores, and periods for identifiers
+        # This matches Athena's identifier naming rules
+        if not _IDENTIFIER_PATTERN.match(identifier):
+            raise ValueError(
+                f"Identifier '{identifier}' contains unsafe characters. Only alphanumeric characters, underscores, and periods are allowed."
+            )
+
+        return identifier
+
     @classmethod
     def _casted_partition_key(cls, key: str) -> str:
         # We need to cast the partition keys to a VARCHAR, since otherwise
         # Athena may throw an error during concatenation / comparison.
-        return f"CAST({key} as VARCHAR)"
+        sanitized_key = cls._sanitize_identifier(key)
+        return f"CAST({sanitized_key} as VARCHAR)"
+
+    @classmethod
+    def _build_max_partition_query(
+        cls, schema: str, table: str, partitions: List[str]
+    ) -> str:
+        """Build SQL query to find the row with maximum partition values.
+
+        Args:
+            schema: Database schema name
+            table: Table name
+            partitions: List of partition column names
+
+        Returns:
+            SQL query string to find the maximum partition
+
+        Raises:
+            ValueError: If any identifier contains unsafe characters
+        """
+        # Sanitize all identifiers to prevent SQL injection
+        sanitized_schema = cls._sanitize_identifier(schema)
+        sanitized_table = cls._sanitize_identifier(table)
+        sanitized_partitions = [
+            cls._sanitize_identifier(partition) for partition in partitions
+        ]
+
+        casted_keys = [cls._casted_partition_key(key) for key in partitions]
+        if len(casted_keys) == 1:
+            part_concat = casted_keys[0]
+        else:
+            separator = "CAST('-' AS VARCHAR)"
+            part_concat = f"CONCAT({f', {separator}, '.join(casted_keys)})"
+
+        return f'select {",".join(sanitized_partitions)} from "{sanitized_schema}"."{sanitized_table}$partitions" where {part_concat} = (select max({part_concat}) from "{sanitized_schema}"."{sanitized_table}$partitions")'
 
     @override
     def get_partitions(
         self, inspector: Inspector, schema: str, table: str
     ) -> Optional[List[str]]:
-        if (
-            not self.config.extract_partitions
-            and not self.config.extract_partitions_using_create_statements
-        ):
+        if not self.config.extract_partitions:
             return None
 
         if not self.cursor:
@@ -557,11 +618,9 @@ class AthenaSource(SQLAlchemySource):
             context=f"{schema}.{table}",
             level=StructuredLogLevel.WARN,
         ):
-            # We create an artifical concatenated partition key to be able to query max partition easier
-            part_concat = " || '-' || ".join(
-                self._casted_partition_key(key) for key in partitions
+            max_partition_query = self._build_max_partition_query(
+                schema, table, partitions
             )
-            max_partition_query = f'select {",".join(partitions)} from "{schema}"."{table}$partitions" where {part_concat} = (select max({part_concat}) from "{schema}"."{table}$partitions")'
             ret = self.cursor.execute(max_partition_query)
             max_partition: Dict[str, str] = {}
             if ret:
@@ -678,16 +737,34 @@ class AthenaSource(SQLAlchemySource):
         ).get(table, None)
 
         if partition and partition.max_partition:
-            max_partition_filters = []
-            for key, value in partition.max_partition.items():
-                max_partition_filters.append(
-                    f"{self._casted_partition_key(key)} = '{value}'"
+            try:
+                # Sanitize identifiers to prevent SQL injection
+                sanitized_schema = self._sanitize_identifier(schema)
+                sanitized_table = self._sanitize_identifier(table)
+
+                max_partition_filters = []
+                for key, value in partition.max_partition.items():
+                    # Sanitize partition key and properly escape the value
+                    sanitized_key = self._sanitize_identifier(key)
+                    # Escape single quotes in the value to prevent injection
+                    escaped_value = value.replace("'", "''") if value else ""
+                    max_partition_filters.append(
+                        f"{self._casted_partition_key(sanitized_key)} = '{escaped_value}'"
+                    )
+                max_partition = str(partition.max_partition)
+                return (
+                    max_partition,
+                    f'SELECT * FROM "{sanitized_schema}"."{sanitized_table}" WHERE {" AND ".join(max_partition_filters)}',
                 )
-            max_partition = str(partition.max_partition)
-            return (
-                max_partition,
-                f'SELECT * FROM "{schema}"."{table}" WHERE {" AND ".join(max_partition_filters)}',
-            )
+            except ValueError as e:
+                # If sanitization fails due to malicious identifiers,
+                # return None to disable partition profiling for this table
+                # rather than crashing the entire ingestion
+                logger.warning(
+                    f"Failed to generate partition profiler query for {schema}.{table} due to unsafe identifiers: {e}. "
+                    f"Partition profiling disabled for this table."
+                )
+                return None, None
         return None, None
 
     def close(self):

datahub/ingestion/source/sql/athena_properties_extractor.py

@@ -174,20 +174,16 @@ class AthenaPropertiesExtractor:
     def format_column_definition(line):
         # Use regex to parse the line more accurately
         # Pattern: column_name data_type [COMMENT comment_text] [,]
-        # Use greedy match for comment to capture everything until trailing comma
-        pattern = r"^\s*(.+?)\s+([\s,\w<>\[\]]+)((\s+COMMENT\s+(.+?)(,?))|(,?)\s*)?$"
-        match = re.match(pattern, line, re.IGNORECASE)
+        # Improved pattern to better separate column name, data type, and comment
+        pattern = r"^\s*([`\w']+)\s+([\w<>\[\](),\s]+?)(\s+COMMENT\s+(.+?))?(,?)\s*$"
+        match = re.match(pattern, line.strip(), re.IGNORECASE)
 
         if not match:
             return line
-        column_name = match.group(1)
-        data_type = match.group(2)
-        comment_part = match.group(5)  # COMMENT part
-        # there are different number of match groups depending on whether comment exists
-        if comment_part:
-            trailing_comma = match.group(6) if match.group(6) else ""
-        else:
-            trailing_comma = match.group(7) if match.group(7) else ""
+        column_name = match.group(1).strip()
+        data_type = match.group(2).strip()
+        comment_part = match.group(4)  # COMMENT part
+        trailing_comma = match.group(5) if match.group(5) else ""
 
         # Add backticks to column name if not already present
         if not (column_name.startswith("`") and column_name.endswith("`")):
@@ -201,17 +197,19 @@ class AthenaPropertiesExtractor:
 
             # Handle comment quoting and escaping
             if comment_part.startswith("'") and comment_part.endswith("'"):
-                # Already properly single quoted - keep as is
-                formatted_comment = comment_part
+                # Already single quoted - but check for proper escaping
+                inner_content = comment_part[1:-1]
+                # Re-escape any single quotes that aren't properly escaped
+                escaped_content = inner_content.replace("'", "''")
+                formatted_comment = f"'{escaped_content}'"
             elif comment_part.startswith('"') and comment_part.endswith('"'):
                 # Double quoted - convert to single quotes and escape internal single quotes
                 inner_content = comment_part[1:-1]
                 escaped_content = inner_content.replace("'", "''")
                 formatted_comment = f"'{escaped_content}'"
             else:
-                # Not quoted - add quotes and escape any single quotes
-                escaped_content = comment_part.replace("'", "''")
-                formatted_comment = f"'{escaped_content}'"
+                # Not quoted - use double quotes to avoid escaping issues with single quotes
+                formatted_comment = f'"{comment_part}"'
 
             result_parts.extend(["COMMENT", formatted_comment])
 
@@ -240,19 +238,39 @@ class AthenaPropertiesExtractor:
                 formatted_lines.append(line)
                 continue
 
-            # Check if we're exiting column definitions (closing parenthesis before PARTITIONED BY or end)
-            if in_column_definition and ")" in line:
-                in_column_definition = False
+            # Skip processing PARTITIONED BY clauses as column definitions
+            if in_column_definition and "PARTITIONED BY" in line.upper():
                 formatted_lines.append(line)
                 continue
 
-            # Process only column definitions (not PARTITIONED BY or other sections)
+            # Process column definitions first, then check for exit condition
             if in_column_definition and stripped_line:
-                # Match column definition pattern and format it
-                formatted_line = AthenaPropertiesExtractor.format_column_definition(
-                    line
-                )
-                formatted_lines.append(formatted_line)
+                # Check if this line contains a column definition (before the closing paren)
+                if ")" in line:
+                    # Split the line at the closing parenthesis
+                    paren_index = line.find(")")
+                    column_part = line[:paren_index].strip()
+                    closing_part = line[paren_index:]
+
+                    if column_part:
+                        # Format the column part
+                        formatted_column = (
+                            AthenaPropertiesExtractor.format_column_definition(
+                                column_part
+                            )
+                        )
+                        # Reconstruct the line
+                        formatted_line = formatted_column.rstrip() + closing_part
+                        formatted_lines.append(formatted_line)
+                    else:
+                        formatted_lines.append(line)
+                    in_column_definition = False
+                else:
+                    # Regular column definition line
+                    formatted_line = AthenaPropertiesExtractor.format_column_definition(
+                        line
+                    )
+                    formatted_lines.append(formatted_line)
             else:
                 # For all other lines, keep as-is
                 formatted_lines.append(line)

datahub/ingestion/source/superset.py

@@ -154,6 +154,7 @@ class SupersetDataset(BaseModel):
     table_name: str
     changed_on_utc: Optional[str] = None
     explore_url: Optional[str] = ""
+    description: Optional[str] = ""
 
     @property
     def modified_dt(self) -> Optional[datetime]:
@@ -1062,7 +1063,7 @@ class SupersetSource(StatefulIngestionSourceBase):
                 fieldPath=col.get("column_name", ""),
                 type=SchemaFieldDataType(data_type),
                 nativeDataType="",
-                description=col.get("column_name", ""),
+                description=col.get("description") or col.get("column_name", ""),
                 nullable=True,
             )
             schema_fields.append(field)
@@ -1283,7 +1284,7 @@ class SupersetSource(StatefulIngestionSourceBase):
 
         dataset_info = DatasetPropertiesClass(
             name=dataset.table_name,
-            description="",
+            description=dataset.description or "",
             externalUrl=dataset_url,
             lastModified=TimeStamp(time=modified_ts),
         )

datahub/ingestion/source/tableau/tableau.py

@@ -1561,12 +1561,15 @@ class TableauSiteSource:
                         }}""",
                     )
             else:
-                # As of Tableau Server 2024.2, the metadata API sporadically returns a 30-second
-                # timeout error.
-                # It doesn't reliably happen, so retrying a couple of times makes sense.
                 if all(
+                    # As of Tableau Server 2024.2, the metadata API sporadically returns a 30-second
+                    # timeout error.
+                    # It doesn't reliably happen, so retrying a couple of times makes sense.
                     error.get("message")
                     == "Execution canceled because timeout of 30000 millis was reached"
+                    # The Metadata API sometimes returns an 'unexpected error' message when querying
+                    # embeddedDatasourcesConnection. Try retrying a couple of times.
+                    or error.get("message") == "Unexpected error occurred"
                     for error in errors
                 ):
                     # If it was only a timeout error, we can retry.
@@ -1578,8 +1581,8 @@ class TableauSiteSource:
                         (self.config.max_retries - retries_remaining + 1) ** 2, 60
                     )
                     logger.info(
-                        f"Query {connection_type} received a 30 second timeout error - will retry in {backoff_time} seconds. "
-                        f"Retries remaining: {retries_remaining}"
+                        f"Query {connection_type} received a retryable error with {retries_remaining} retries remaining, "
+                        f"will retry in {backoff_time} seconds: {errors}"
                     )
                     time.sleep(backoff_time)
                     return self.get_connection_object_page(

datahub/ingestion/source/unity/config.py

@@ -35,6 +35,10 @@ from datahub.utilities.global_warning_util import add_global_warning
 
 logger = logging.getLogger(__name__)
 
+# Configuration default constants
+INCLUDE_TAGS_DEFAULT = True
+INCLUDE_HIVE_METASTORE_DEFAULT = True
+
 
 class LineageDataSource(ConfigEnum):
     AUTO = "AUTO"
@@ -137,10 +141,18 @@ class UnityCatalogSourceConfig(
     )
     warehouse_id: Optional[str] = pydantic.Field(
         default=None,
-        description="SQL Warehouse id, for running queries. If not set, will use the default warehouse.",
+        description=(
+            "SQL Warehouse id, for running queries. Must be explicitly provided to enable SQL-based features. "
+            "Required for the following features that need SQL access: "
+            "1) Tag extraction (include_tags=True) - queries system.information_schema.tags "
+            "2) Hive Metastore catalog (include_hive_metastore=True) - queries legacy hive_metastore catalog "
+            "3) System table lineage (lineage_data_source=SYSTEM_TABLES) - queries system.access.table_lineage/column_lineage "
+            "4) Data profiling (profiling.enabled=True) - runs SELECT/ANALYZE queries on tables. "
+            "When warehouse_id is missing, these features will be automatically disabled (with warnings) to allow ingestion to continue."
+        ),
     )
     include_hive_metastore: bool = pydantic.Field(
-        default=True,
+        default=INCLUDE_HIVE_METASTORE_DEFAULT,
         description="Whether to ingest legacy `hive_metastore` catalog. This requires executing queries on SQL warehouse.",
     )
     workspace_name: Optional[str] = pydantic.Field(
@@ -236,8 +248,12 @@ class UnityCatalogSourceConfig(
     )
 
     include_tags: bool = pydantic.Field(
-        default=True,
-        description="Option to enable/disable column/table tag extraction.",
+        default=INCLUDE_TAGS_DEFAULT,
+        description=(
+            "Option to enable/disable column/table tag extraction. "
+            "Requires warehouse_id to be set since tag extraction needs to query system.information_schema.tags. "
+            "If warehouse_id is not provided, this will be automatically disabled to allow ingestion to continue."
+        ),
     )
 
     _rename_table_ownership = pydantic_renamed_field(
@@ -310,8 +326,6 @@ class UnityCatalogSourceConfig(
         description="Details about the delta lake, incase to emit siblings",
     )
 
-    scheme: str = DATABRICKS
-
     include_ml_model_aliases: bool = pydantic.Field(
         default=False,
         description="Whether to include ML model aliases in the ingestion.",
@@ -323,6 +337,51 @@ class UnityCatalogSourceConfig(
         description="Maximum number of ML models to ingest.",
     )
 
+    _forced_disable_tag_extraction: bool = pydantic.PrivateAttr(default=False)
+    _forced_disable_hive_metastore_extraction = pydantic.PrivateAttr(default=False)
+
+    scheme: str = DATABRICKS
+
+    def __init__(self, **data):
+        # First, let the parent handle the root validators and field processing
+        super().__init__(**data)
+
+        # After model creation, check if we need to auto-disable features
+        # based on the final warehouse_id value (which may have been set by root validators)
+        include_tags_original = data.get("include_tags", INCLUDE_TAGS_DEFAULT)
+        include_hive_metastore_original = data.get(
+            "include_hive_metastore", INCLUDE_HIVE_METASTORE_DEFAULT
+        )
+
+        # Track what we're force-disabling
+        forced_disable_tag_extraction = False
+        forced_disable_hive_metastore_extraction = False
+
+        # Check if features should be auto-disabled based on final warehouse_id
+        if include_tags_original and not self.warehouse_id:
+            forced_disable_tag_extraction = True
+            self.include_tags = False  # Modify the model attribute directly
+            logger.warning(
+                "warehouse_id is not set but include_tags=True. "
+                "Automatically disabling tag extraction since it requires SQL queries. "
+                "Set warehouse_id to enable tag extraction."
+            )
+
+        if include_hive_metastore_original and not self.warehouse_id:
+            forced_disable_hive_metastore_extraction = True
+            self.include_hive_metastore = False  # Modify the model attribute directly
+            logger.warning(
+                "warehouse_id is not set but include_hive_metastore=True. "
+                "Automatically disabling hive metastore extraction since it requires SQL queries. "
+                "Set warehouse_id to enable hive metastore extraction."
+            )
+
+        # Set private attributes
+        self._forced_disable_tag_extraction = forced_disable_tag_extraction
+        self._forced_disable_hive_metastore_extraction = (
+            forced_disable_hive_metastore_extraction
+        )
+
     def get_sql_alchemy_url(self, database: Optional[str] = None) -> str:
         uri_opts = {"http_path": f"/sql/1.0/warehouses/{self.warehouse_id}"}
         if database:
@@ -392,11 +451,6 @@ class UnityCatalogSourceConfig(
                 "When `warehouse_id` is set, it must match the `warehouse_id` in `profiling`."
             )
 
-        if values.get("include_hive_metastore") and not values.get("warehouse_id"):
-            raise ValueError(
-                "When `include_hive_metastore` is set, `warehouse_id` must be set."
-            )
-
         if values.get("warehouse_id") and profiling and not profiling.warehouse_id:
             profiling.warehouse_id = values["warehouse_id"]
 

datahub/ingestion/source/unity/proxy.py

@@ -4,6 +4,7 @@ Manage the communication with DataBricks Server and provide equivalent dataclass
 
 import dataclasses
 import logging
+import os
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Union, cast
@@ -71,6 +72,23 @@ logger: logging.Logger = logging.getLogger(__name__)
 _MAX_CONCURRENT_CATALOGS = 1
 
 
+# Import and apply the proxy patch from separate module
+try:
+    from datahub.ingestion.source.unity.proxy_patch import (
+        apply_databricks_proxy_fix,
+        mask_proxy_credentials,
+    )
+
+    # Apply the fix when the module is imported
+    apply_databricks_proxy_fix()
+except ImportError as e:
+    logger.debug(f"Could not import proxy patch module: {e}")
+
+    # Fallback function for masking credentials
+    def mask_proxy_credentials(url: Optional[str]) -> str:
+        return "***MASKED***" if url else "None"
+
+
 @dataclasses.dataclass
 class TableInfoWithGeneration(TableInfo):
     generation: Optional[int] = None
@@ -411,7 +429,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
             query = f"""
                 SELECT
                     entity_type, entity_id,
-                    source_table_full_name, source_type,
+                    source_table_full_name, source_type, source_path,
                     target_table_full_name, target_type,
                     max(event_time) as last_updated
                 FROM system.access.table_lineage
@@ -420,7 +438,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     {additional_where}
                 GROUP BY
                     entity_type, entity_id,
-                    source_table_full_name, source_type,
+                    source_table_full_name, source_type, source_path,
                     target_table_full_name, target_type
                 """
             rows = self._execute_sql_query(query, [catalog, catalog])
@@ -432,6 +450,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                 source_full_name = row["source_table_full_name"]
                 target_full_name = row["target_table_full_name"]
                 source_type = row["source_type"]
+                source_path = row["source_path"]
                 last_updated = row["last_updated"]
 
                 # Initialize TableLineageInfo for both source and target tables if they're in our catalog
@@ -460,7 +479,7 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
                     # Handle external upstreams (PATH type)
                     elif source_type == "PATH":
                         external_upstream = ExternalUpstream(
-                            path=source_full_name,
+                            path=source_path,
                             source_type=source_type,
                             last_updated=last_updated,
                         )
@@ -973,16 +992,82 @@ class UnityCatalogApiProxy(UnityCatalogProxyProfilingMixin):
 
     def _execute_sql_query(self, query: str, params: Sequence[Any] = ()) -> List[Row]:
         """Execute SQL query using databricks-sql connector for better performance"""
+        logger.debug(f"Executing SQL query with {len(params)} parameters")
+        if logger.isEnabledFor(logging.DEBUG):
+            # Only log full query in debug mode to avoid performance overhead
+            logger.debug(f"Full SQL query: {query}")
+            if params:
+                logger.debug(f"Query parameters: {params}")
+
+        # Check if warehouse_id is available for SQL operations
+        if not self.warehouse_id:
+            self.report.report_warning(
+                "Cannot execute SQL query",
+                "warehouse_id is not configured. SQL operations require a valid warehouse_id to be set in the Unity Catalog configuration",
+            )
+            logger.warning(
+                "Cannot execute SQL query: warehouse_id is not configured. "
+                "SQL operations require a valid warehouse_id to be set in the Unity Catalog configuration."
+            )
+            return []
+
+        # Log connection parameters (with masked token)
+        masked_params = {**self._sql_connection_params}
+        if "access_token" in masked_params:
+            masked_params["access_token"] = "***MASKED***"
+        logger.debug(f"Using connection parameters: {masked_params}")
+
+        # Log proxy environment variables that affect SQL connections
+        proxy_env_debug = {}
+        for var in ["HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy"]:
+            value = os.environ.get(var)
+            if value:
+                proxy_env_debug[var] = mask_proxy_credentials(value)
+
+        if proxy_env_debug:
+            logger.debug(
+                f"SQL connection will use proxy environment variables: {proxy_env_debug}"
+            )
+        else:
+            logger.debug("No proxy environment variables detected for SQL connection")
+
         try:
             with (
                 connect(**self._sql_connection_params) as connection,
                 connection.cursor() as cursor,
             ):
                 cursor.execute(query, list(params))
-                return cursor.fetchall()
+                rows = cursor.fetchall()
+                logger.debug(
+                    f"SQL query executed successfully, returned {len(rows)} rows"
+                )
+                return rows
 
         except Exception as e:
-            logger.warning(f"Failed to execute SQL query: {e}")
+            logger.warning(f"Failed to execute SQL query: {e}", exc_info=True)
+            if logger.isEnabledFor(logging.DEBUG):
+                # Only log failed query details in debug mode for security
+                logger.debug(f"SQL query that failed: {query}")
+                logger.debug(f"SQL query parameters: {params}")
+
+            # Check if this might be a proxy-related error
+            error_str = str(e).lower()
+            if any(
+                proxy_keyword in error_str
+                for proxy_keyword in [
+                    "proxy",
+                    "407",
+                    "authentication required",
+                    "tunnel",
+                    "connect",
+                ]
+            ):
+                logger.error(
+                    "SQL query failure appears to be proxy-related. "
+                    "Please check proxy configuration and authentication. "
+                    f"Proxy environment variables detected: {list(proxy_env_debug.keys())}"
+                )
+
             return []
 
     @cached(cachetools.FIFOCache(maxsize=_MAX_CONCURRENT_CATALOGS))