Merge branch 'main' into sea-http-client

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/backend.py

@@ -19,6 +19,7 @@
     WaitTimeout,
     MetadataCommands,
 )
+from databricks.sql.backend.sea.utils.normalize import normalize_sea_type_to_thrift
 from databricks.sql.thrift_api.TCLIService import ttypes
 
 if TYPE_CHECKING:
@@ -322,6 +323,11 @@ def _extract_description_from_manifest(
             # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
             name = col_data.get("name", "")
             type_name = col_data.get("type_name", "")
+
+            # Normalize SEA type to Thrift conventions before any processing
+            type_name = normalize_sea_type_to_thrift(type_name, col_data)
+
+            # Now strip _TYPE suffix and convert to lowercase
             type_name = (
                 type_name[:-5] if type_name.endswith("_TYPE") else type_name
             ).lower()

src/databricks/sql/backend/sea/result_set.py

@@ -92,20 +92,19 @@ def _convert_json_types(self, row: List[str]) -> List[Any]:
         converted_row = []
 
         for i, value in enumerate(row):
+            column_name = self.description[i][0]
             column_type = self.description[i][1]
             precision = self.description[i][4]
             scale = self.description[i][5]
 
-            try:
-                converted_value = SqlTypeConverter.convert_value(
-                    value, column_type, precision=precision, scale=scale
-                )
-                converted_row.append(converted_value)
-            except Exception as e:
-                logger.warning(
-                    f"Error converting value '{value}' to {column_type}: {e}"
-                )
-                converted_row.append(value)
+            converted_value = SqlTypeConverter.convert_value(
+                value,
+                column_type,
+                column_name=column_name,
+                precision=precision,
+                scale=scale,
+            )
+            converted_row.append(converted_value)
 
         return converted_row
 

src/databricks/sql/backend/sea/utils/conversion.py

@@ -50,60 +50,65 @@ def _convert_decimal(
 
 class SqlType:
     """
-    SQL type constants
+    SQL type constants based on Thrift TTypeId values.
 
-    The list of types can be found in the SEA REST API Reference:
-    https://docs.databricks.com/api/workspace/statementexecution/executestatement
+    These correspond to the normalized type names that come from the SEA backend
+    after normalize_sea_type_to_thrift processing (lowercase, without _TYPE suffix).
     """
 
     # Numeric types
-    BYTE = "byte"
-    SHORT = "short"
-    INT = "int"
-    LONG = "long"
-    FLOAT = "float"
-    DOUBLE = "double"
-    DECIMAL = "decimal"
+    TINYINT = "tinyint"  # Maps to TTypeId.TINYINT_TYPE
+    SMALLINT = "smallint"  # Maps to TTypeId.SMALLINT_TYPE
+    INT = "int"  # Maps to TTypeId.INT_TYPE
+    BIGINT = "bigint"  # Maps to TTypeId.BIGINT_TYPE
+    FLOAT = "float"  # Maps to TTypeId.FLOAT_TYPE
+    DOUBLE = "double"  # Maps to TTypeId.DOUBLE_TYPE
+    DECIMAL = "decimal"  # Maps to TTypeId.DECIMAL_TYPE
 
     # Boolean type
-    BOOLEAN = "boolean"
+    BOOLEAN = "boolean"  # Maps to TTypeId.BOOLEAN_TYPE
 
     # Date/Time types
-    DATE = "date"
-    TIMESTAMP = "timestamp"
-    INTERVAL = "interval"
+    DATE = "date"  # Maps to TTypeId.DATE_TYPE
+    TIMESTAMP = "timestamp"  # Maps to TTypeId.TIMESTAMP_TYPE
+    INTERVAL_YEAR_MONTH = (
+        "interval_year_month"  # Maps to TTypeId.INTERVAL_YEAR_MONTH_TYPE
+    )
+    INTERVAL_DAY_TIME = "interval_day_time"  # Maps to TTypeId.INTERVAL_DAY_TIME_TYPE
 
     # String types
-    CHAR = "char"
-    STRING = "string"
+    CHAR = "char"  # Maps to TTypeId.CHAR_TYPE
+    VARCHAR = "varchar"  # Maps to TTypeId.VARCHAR_TYPE
+    STRING = "string"  # Maps to TTypeId.STRING_TYPE
 
     # Binary type
-    BINARY = "binary"
+    BINARY = "binary"  # Maps to TTypeId.BINARY_TYPE
 
     # Complex types
-    ARRAY = "array"
-    MAP = "map"
-    STRUCT = "struct"
+    ARRAY = "array"  # Maps to TTypeId.ARRAY_TYPE
+    MAP = "map"  # Maps to TTypeId.MAP_TYPE
+    STRUCT = "struct"  # Maps to TTypeId.STRUCT_TYPE
 
     # Other types
-    NULL = "null"
-    USER_DEFINED_TYPE = "user_defined_type"
+    NULL = "null"  # Maps to TTypeId.NULL_TYPE
+    UNION = "union"  # Maps to TTypeId.UNION_TYPE
+    USER_DEFINED = "user_defined"  # Maps to TTypeId.USER_DEFINED_TYPE
 
 
 class SqlTypeConverter:
     """
     Utility class for converting SQL types to Python types.
-    Based on the types supported by the Databricks SDK.
+    Based on the Thrift TTypeId types after normalization.
     """
 
     # SQL type to conversion function mapping
     # TODO: complex types
     TYPE_MAPPING: Dict[str, Callable] = {
         # Numeric types
-        SqlType.BYTE: lambda v: int(v),
-        SqlType.SHORT: lambda v: int(v),
+        SqlType.TINYINT: lambda v: int(v),
+        SqlType.SMALLINT: lambda v: int(v),
         SqlType.INT: lambda v: int(v),
-        SqlType.LONG: lambda v: int(v),
+        SqlType.BIGINT: lambda v: int(v),
         SqlType.FLOAT: lambda v: float(v),
         SqlType.DOUBLE: lambda v: float(v),
         SqlType.DECIMAL: _convert_decimal,
@@ -112,30 +117,34 @@ class SqlTypeConverter:
         # Date/Time types
         SqlType.DATE: lambda v: datetime.date.fromisoformat(v),
         SqlType.TIMESTAMP: lambda v: parser.parse(v),
-        SqlType.INTERVAL: lambda v: v,  # Keep as string for now
+        SqlType.INTERVAL_YEAR_MONTH: lambda v: v,  # Keep as string for now
+        SqlType.INTERVAL_DAY_TIME: lambda v: v,  # Keep as string for now
         # String types - no conversion needed
         SqlType.CHAR: lambda v: v,
+        SqlType.VARCHAR: lambda v: v,
         SqlType.STRING: lambda v: v,
         # Binary type
         SqlType.BINARY: lambda v: bytes.fromhex(v),
         # Other types
         SqlType.NULL: lambda v: None,
         # Complex types and user-defined types return as-is
-        SqlType.USER_DEFINED_TYPE: lambda v: v,
+        SqlType.USER_DEFINED: lambda v: v,
     }
 
     @staticmethod
     def convert_value(
         value: str,
         sql_type: str,
+        column_name: Optional[str],
         **kwargs,
     ) -> object:
         """
         Convert a string value to the appropriate Python type based on SQL type.
 
         Args:
             value: The string value to convert
-            sql_type: The SQL type (e.g., 'int', 'decimal')
+            sql_type: The SQL type (e.g., 'tinyint', 'decimal')
+            column_name: The name of the column being converted
             **kwargs: Additional keyword arguments for the conversion function
 
         Returns:
@@ -155,6 +164,10 @@ def convert_value(
                 return converter_func(value, precision, scale)
             else:
                 return converter_func(value)
-        except (ValueError, TypeError, decimal.InvalidOperation) as e:
-            logger.warning(f"Error converting value '{value}' to {sql_type}: {e}")
+        except Exception as e:
+            warning_message = f"Error converting value '{value}' to {sql_type}"
+            if column_name:
+                warning_message += f" in column {column_name}"
+            warning_message += f": {e}"
+            logger.warning(warning_message)
             return value

src/databricks/sql/backend/sea/utils/normalize.py

@@ -0,0 +1,50 @@
+"""
+Type normalization utilities for SEA backend.
+
+This module provides functionality to normalize SEA type names to match
+Thrift type naming conventions.
+"""
+
+from typing import Dict, Any
+
+# SEA types that need to be translated to Thrift types
+# The list of all SEA types is available in the REST reference at:
+#    https://docs.databricks.com/api/workspace/statementexecution/executestatement
+# The list of all Thrift types can be found in the ttypes.TTypeId definition
+# The SEA types that do not align with Thrift are explicitly mapped below
+SEA_TO_THRIFT_TYPE_MAP = {
+    "BYTE": "TINYINT",
+    "SHORT": "SMALLINT",
+    "LONG": "BIGINT",
+    "INTERVAL": "INTERVAL",  # Default mapping, will be overridden if type_interval_type is present
+}
+
+
+def normalize_sea_type_to_thrift(type_name: str, col_data: Dict[str, Any]) -> str:
+    """
+    Normalize SEA type names to match Thrift type naming conventions.
+
+    Args:
+        type_name: The type name from SEA (e.g., "BYTE", "LONG", "INTERVAL")
+        col_data: The full column data dictionary from manifest (for accessing type_interval_type)
+
+    Returns:
+        Normalized type name matching Thrift conventions
+    """
+    # Early return if type doesn't need mapping
+    if type_name not in SEA_TO_THRIFT_TYPE_MAP:
+        return type_name
+
+    normalized_type = SEA_TO_THRIFT_TYPE_MAP[type_name]
+
+    # Special handling for interval types
+    if type_name == "INTERVAL":
+        type_interval_type = col_data.get("type_interval_type")
+        if type_interval_type:
+            return (
+                "INTERVAL_YEAR_MONTH"
+                if any(t in type_interval_type.upper() for t in ["YEAR", "MONTH"])
+                else "INTERVAL_DAY_TIME"
+            )
+
+    return normalized_type

src/databricks/sql/client.py

@@ -248,11 +248,8 @@ def read(self) -> Optional[OAuthToken]:
         self.lz4_compression = kwargs.get("enable_query_result_lz4_compression", True)
         self.use_cloud_fetch = kwargs.get("use_cloud_fetch", True)
         self._cursors = []  # type: List[Cursor]
-
-        self.server_telemetry_enabled = True
-        self.client_telemetry_enabled = kwargs.get("enable_telemetry", False)
-        self.telemetry_enabled = (
-            self.client_telemetry_enabled and self.server_telemetry_enabled
+        self.telemetry_batch_size = kwargs.get(
+            "telemetry_batch_size", TelemetryClientFactory.DEFAULT_BATCH_SIZE
         )
 
         try:
@@ -285,11 +282,16 @@ def read(self) -> Optional[OAuthToken]:
         )
         self.staging_allowed_local_path = kwargs.get("staging_allowed_local_path", None)
 
+        self.force_enable_telemetry = kwargs.get("force_enable_telemetry", False)
+        self.enable_telemetry = kwargs.get("enable_telemetry", False)
+        self.telemetry_enabled = TelemetryHelper.is_telemetry_enabled(self)
+
         TelemetryClientFactory.initialize_telemetry_client(
             telemetry_enabled=self.telemetry_enabled,
             session_id_hex=self.get_session_id_hex(),
             auth_provider=self.session.auth_provider,
             host_url=self.session.host,
+            batch_size=self.telemetry_batch_size,
         )
 
         self._telemetry_client = TelemetryClientFactory.get_telemetry_client(

src/databricks/sql/cloudfetch/downloader.py

@@ -54,12 +54,14 @@ class DownloadableResultSettings:
         link_expiry_buffer_secs (int): Time in seconds to prevent download of a link before it expires. Default 0 secs.
         download_timeout (int): Timeout for download requests. Default 60 secs.
         max_consecutive_file_download_retries (int): Number of consecutive download retries before shutting down.
+        min_cloudfetch_download_speed (float): Threshold in MB/s below which to log warning. Default 0.1 MB/s.
     """
 
     is_lz4_compressed: bool
     link_expiry_buffer_secs: int = 0
     download_timeout: int = 60
     max_consecutive_file_download_retries: int = 0
+    min_cloudfetch_download_speed: float = 0.1
 
 
 class ResultSetDownloadHandler:
@@ -100,6 +102,8 @@ def run(self) -> DownloadedFile:
             self.link, self.settings.link_expiry_buffer_secs
         )
 
+        start_time = time.time()
+
         with self._http_client.execute(
             method=HttpMethod.GET,
             url=self.link.fileLink,
@@ -112,6 +116,13 @@ def run(self) -> DownloadedFile:
 
             # Save (and decompress if needed) the downloaded file
             compressed_data = response.content
+
+            # Log download metrics
+            download_duration = time.time() - start_time
+            self._log_download_metrics(
+                self.link.fileLink, len(compressed_data), download_duration
+            )
+
             decompressed_data = (
                 ResultSetDownloadHandler._decompress_data(compressed_data)
                 if self.settings.is_lz4_compressed
@@ -138,6 +149,32 @@ def run(self) -> DownloadedFile:
                 self.link.rowCount,
             )
 
+    def _log_download_metrics(
+        self, url: str, bytes_downloaded: int, duration_seconds: float
+    ):
+        """Log download speed metrics at INFO/WARN levels."""
+        # Calculate speed in MB/s (ensure float division for precision)
+        speed_mbps = (float(bytes_downloaded) / (1024 * 1024)) / duration_seconds
+
+        urlEndpoint = url.split("?")[0]
+        # INFO level logging
+        logger.info(
+            "CloudFetch download completed: %.4f MB/s, %d bytes in %.3fs from %s",
+            speed_mbps,
+            bytes_downloaded,
+            duration_seconds,
+            urlEndpoint,
+        )
+
+        # WARN level logging if below threshold
+        if speed_mbps < self.settings.min_cloudfetch_download_speed:
+            logger.warning(
+                "CloudFetch download slower than threshold: %.4f MB/s (threshold: %.1f MB/s) from %s",
+                speed_mbps,
+                self.settings.min_cloudfetch_download_speed,
+                url,
+            )
+
     @staticmethod
     def _validate_link(link: TSparkArrowResultLink, expiry_buffer_secs: int):
         """