Merge branch 'main' into telemetry-server-flag

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

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/telemetry/telemetry_client.py

@@ -145,6 +145,9 @@ def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
     def close(self):
         pass
 
+    def _flush(self):
+        pass
+
 
 class TelemetryClient(BaseTelemetryClient):
     """

tests/e2e/test_complex_types.py

@@ -39,9 +39,11 @@ def table_fixture(self, connection_details):
                 )
                 """
             )
-            yield
-            # Clean up the table after the test
-            cursor.execute("DELETE FROM pysql_test_complex_types_table")
+            try:
+                yield
+            finally:
+                # Clean up the table after the test
+                cursor.execute("DELETE FROM pysql_test_complex_types_table")
 
     @pytest.mark.parametrize(
         "field,expected_type",