Merge branch 'fetch-json-inline' into ext-links-sea

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -3,7 +3,13 @@
 
 This script runs all the individual test modules and displays
 a summary of test results with visual indicators.
+
+In order to run the script, the following environment variables need to be set:
+- DATABRICKS_SERVER_HOSTNAME: The hostname of the Databricks server
+- DATABRICKS_HTTP_PATH: The HTTP path of the Databricks server
+- DATABRICKS_TOKEN: The token to use for authentication
 """
+
 import os
 import sys
 import logging

examples/experimental/tests/test_sea_async_query.py

@@ -17,7 +17,7 @@ def test_sea_async_query_with_cloud_fetch():
     Test executing a query asynchronously using the SEA backend with cloud fetch enabled.
 
     This function connects to a Databricks SQL endpoint using the SEA backend,
-    executes a query asynchronously with cloud fetch enabled, and verifies that execution completes successfully.
+    executes a simple query asynchronously with cloud fetch enabled, and verifies that execution completes successfully.
     """
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
@@ -82,9 +82,6 @@ def test_sea_async_query_with_cloud_fetch():
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
         actual_row_count = len(results)
-        logger.info(
-            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
-        )
 
         logger.info(
             f"Requested {requested_row_count} rows, received {actual_row_count} rows"
@@ -123,7 +120,7 @@ def test_sea_async_query_without_cloud_fetch():
     Test executing a query asynchronously using the SEA backend with cloud fetch disabled.
 
     This function connects to a Databricks SQL endpoint using the SEA backend,
-    executes a query asynchronously with cloud fetch disabled, and verifies that execution completes successfully.
+    executes a simple query asynchronously with cloud fetch disabled, and verifies that execution completes successfully.
     """
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
@@ -188,9 +185,6 @@ def test_sea_async_query_without_cloud_fetch():
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
         actual_row_count = len(results)
-        logger.info(
-            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
-        )
 
         logger.info(
             f"Requested {requested_row_count} rows, received {actual_row_count} rows"

examples/experimental/tests/test_sea_sync_query.py

@@ -6,7 +6,7 @@
 import logging
 from databricks.sql.client import Connection
 
-logging.basicConfig(level=logging.DEBUG)
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
@@ -15,7 +15,7 @@ def test_sea_sync_query_with_cloud_fetch():
     Test executing a query synchronously using the SEA backend with cloud fetch enabled.
 
     This function connects to a Databricks SQL endpoint using the SEA backend,
-    executes a query with cloud fetch enabled, and verifies that execution completes successfully.
+    executes a simple query with cloud fetch enabled, and verifies that execution completes successfully.
     """
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
@@ -93,7 +93,7 @@ def test_sea_sync_query_without_cloud_fetch():
     Test executing a query synchronously using the SEA backend with cloud fetch disabled.
 
     This function connects to a Databricks SQL endpoint using the SEA backend,
-    executes a query with cloud fetch disabled, and verifies that execution completes successfully.
+    executes a simple query with cloud fetch disabled, and verifies that execution completes successfully.
     """
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")

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

@@ -15,8 +15,6 @@
     MetadataCommands,
 )
 
-from databricks.sql.thrift_api.TCLIService import ttypes
-
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
     from databricks.sql.result_set import SeaResultSet
@@ -409,7 +407,7 @@ def execute_command(
         lz4_compression: bool,
         cursor: Cursor,
         use_cloud_fetch: bool,
-        parameters: List[ttypes.TSparkParameter],
+        parameters: List[Dict[str, Any]],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
     ) -> Union[SeaResultSet, None]:
@@ -443,9 +441,9 @@ def execute_command(
             for param in parameters:
                 sea_parameters.append(
                     StatementParameter(
-                        name=param.name,
-                        value=param.value.stringValue,
-                        type=param.type,
+                        name=param["name"],
+                        value=param["value"],
+                        type=param["type"] if "type" in param else None,
                     )
                 )
 
@@ -620,10 +618,10 @@ def get_execution_result(
             connection=cursor.connection,
             execute_response=execute_response,
             sea_client=self,
-            buffer_size_bytes=cursor.buffer_size_bytes,
-            arraysize=cursor.arraysize,
             result_data=response.result,
             manifest=response.manifest,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
         )
 
     def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:

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

@@ -0,0 +1,135 @@
+"""
+Type conversion utilities for the Databricks SQL Connector.
+
+This module provides functionality to convert string values from SEA Inline results
+to appropriate Python types based on column metadata.
+"""
+
+import datetime
+import decimal
+import logging
+from dateutil import parser
+from typing import Any, Callable, Dict, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+class SqlType:
+    """
+    SQL type constants
+
+    The list of types can be found in the SEA REST API Reference:
+    https://docs.databricks.com/api/workspace/statementexecution/executestatement
+    """
+
+    # Numeric types
+    BYTE = "byte"
+    SHORT = "short"
+    INT = "int"
+    LONG = "long"
+    FLOAT = "float"
+    DOUBLE = "double"
+    DECIMAL = "decimal"
+
+    # Boolean type
+    BOOLEAN = "boolean"
+
+    # Date/Time types
+    DATE = "date"
+    TIMESTAMP = "timestamp"
+    INTERVAL = "interval"
+
+    # String types
+    CHAR = "char"
+    STRING = "string"
+
+    # Binary type
+    BINARY = "binary"
+
+    # Complex types
+    ARRAY = "array"
+    MAP = "map"
+    STRUCT = "struct"
+
+    # Other types
+    NULL = "null"
+    USER_DEFINED_TYPE = "user_defined_type"
+
+
+class SqlTypeConverter:
+    """
+    Utility class for converting SQL types to Python types.
+    Based on the types supported by the Databricks SDK.
+    """
+
+    # 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.INT: lambda v: int(v),
+        SqlType.LONG: lambda v: int(v),
+        SqlType.FLOAT: lambda v: float(v),
+        SqlType.DOUBLE: lambda v: float(v),
+        SqlType.DECIMAL: lambda v, p=None, s=None: (
+            decimal.Decimal(v).quantize(
+                decimal.Decimal(f'0.{"0" * s}'), context=decimal.Context(prec=p)
+            )
+            if p is not None and s is not None
+            else decimal.Decimal(v)
+        ),
+        # Boolean type
+        SqlType.BOOLEAN: lambda v: v.lower() in ("true", "t", "1", "yes", "y"),
+        # 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
+        # String types - no conversion needed
+        SqlType.CHAR: 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,
+    }
+
+    @staticmethod
+    def convert_value(
+        value: Any,
+        sql_type: str,
+        precision: Optional[int] = None,
+        scale: Optional[int] = None,
+    ) -> Any:
+        """
+        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')
+            precision: Optional precision for decimal types
+            scale: Optional scale for decimal types
+
+        Returns:
+            The converted value in the appropriate Python type
+        """
+
+        if value is None:
+            return None
+
+        sql_type = sql_type.lower().strip()
+
+        if sql_type not in SqlTypeConverter.TYPE_MAPPING:
+            return value
+
+        converter_func = SqlTypeConverter.TYPE_MAPPING[sql_type]
+        try:
+            if sql_type == SqlType.DECIMAL:
+                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}")
+            return value

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

@@ -77,9 +77,9 @@ def _filter_sea_result_set(
             connection=result_set.connection,
             execute_response=execute_response,
             sea_client=cast(SeaDatabricksClient, result_set.backend),
+            result_data=result_data,
             buffer_size_bytes=result_set.buffer_size_bytes,
             arraysize=result_set.arraysize,
-            result_data=result_data,
         )
 
         return filtered_result_set