Merge branch 'main' into robust-metadata-sea

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:
@@ -323,6 +324,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

@@ -254,6 +254,9 @@ def read(self) -> Optional[OAuthToken]:
         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:
             self.session = Session(
@@ -290,6 +293,7 @@ def read(self) -> Optional[OAuthToken]:
             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/telemetry/telemetry_client.py

@@ -138,19 +138,18 @@ class TelemetryClient(BaseTelemetryClient):
     TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
     TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
 
-    DEFAULT_BATCH_SIZE = 100
-
     def __init__(
         self,
         telemetry_enabled,
         session_id_hex,
         auth_provider,
         host_url,
         executor,
+        batch_size,
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
-        self._batch_size = self.DEFAULT_BATCH_SIZE
+        self._batch_size = batch_size
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None
@@ -318,7 +317,7 @@ def close(self):
 class TelemetryClientFactory:
     """
     Static factory class for creating and managing telemetry clients.
-    It uses a thread pool to handle asynchronous operations.
+    It uses a thread pool to handle asynchronous operations and a single flush thread for all clients.
     """
 
     _clients: Dict[
@@ -331,6 +330,13 @@ class TelemetryClientFactory:
     _original_excepthook = None
     _excepthook_installed = False
 
+    # Shared flush thread for all clients
+    _flush_thread = None
+    _flush_event = threading.Event()
+    _flush_interval_seconds = 90
+
+    DEFAULT_BATCH_SIZE = 100
+
     @classmethod
     def _initialize(cls):
         """Initialize the factory if not already initialized"""
@@ -341,11 +347,39 @@ def _initialize(cls):
                 max_workers=10
             )  # Thread pool for async operations
             cls._install_exception_hook()
+            cls._start_flush_thread()
             cls._initialized = True
             logger.debug(
                 "TelemetryClientFactory initialized with thread pool (max_workers=10)"
             )
 
+    @classmethod
+    def _start_flush_thread(cls):
+        """Start the shared background thread for periodic flushing of all clients"""
+        cls._flush_event.clear()
+        cls._flush_thread = threading.Thread(target=cls._flush_worker, daemon=True)
+        cls._flush_thread.start()
+
+    @classmethod
+    def _flush_worker(cls):
+        """Background worker thread for periodic flushing of all clients"""
+        while not cls._flush_event.wait(cls._flush_interval_seconds):
+            logger.debug("Performing periodic flush for all telemetry clients")
+
+            with cls._lock:
+                clients_to_flush = list(cls._clients.values())
+
+                for client in clients_to_flush:
+                    client._flush()
+
+    @classmethod
+    def _stop_flush_thread(cls):
+        """Stop the shared background flush thread"""
+        if cls._flush_thread is not None:
+            cls._flush_event.set()
+            cls._flush_thread.join(timeout=1.0)
+            cls._flush_thread = None
+
     @classmethod
     def _install_exception_hook(cls):
         """Install global exception handler for unhandled exceptions"""
@@ -374,6 +408,7 @@ def initialize_telemetry_client(
         session_id_hex,
         auth_provider,
         host_url,
+        batch_size,
     ):
         """Initialize a telemetry client for a specific connection if telemetry is enabled"""
         try:
@@ -395,6 +430,7 @@ def initialize_telemetry_client(
                             auth_provider=auth_provider,
                             host_url=host_url,
                             executor=TelemetryClientFactory._executor,
+                            batch_size=batch_size,
                         )
                     else:
                         TelemetryClientFactory._clients[
@@ -433,6 +469,7 @@ def close(session_id_hex):
                     "No more telemetry clients, shutting down thread pool executor"
                 )
                 try:
+                    TelemetryClientFactory._stop_flush_thread()
                     TelemetryClientFactory._executor.shutdown(wait=True)
                     TelemetryHttpClient.close()
                 except Exception as e:
@@ -458,6 +495,7 @@ def connection_failure_log(
             session_id_hex=UNAUTH_DUMMY_SESSION_ID,
             auth_provider=None,
             host_url=host_url,
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
 
         telemetry_client = TelemetryClientFactory.get_telemetry_client(