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

Author: varun-edachali-dbx

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

@@ -6,12 +6,15 @@
 
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
+    ResultFormat,
+    ResultDisposition,
+    ResultCompression,
+    WaitTimeout,
 )
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
     from databricks.sql.result_set import ResultSet
-    from databricks.sql.backend.sea.models.responses import GetChunksResponse
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (
@@ -21,16 +24,10 @@
     BackendType,
     ExecuteResponse,
 )
-from databricks.sql.exc import Error, NotSupportedError, ServerOperationError
+from databricks.sql.exc import ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
-from databricks.sql.utils import SeaResultSetQueueFactory
-from databricks.sql.backend.sea.models.base import (
-    ResultData,
-    ExternalLink,
-    ResultManifest,
-)
 
 from databricks.sql.backend.sea.models import (
     ExecuteStatementRequest,
@@ -45,6 +42,11 @@
     CreateSessionResponse,
     GetChunksResponse,
 )
+from databricks.sql.backend.sea.models.responses import (
+    parse_status,
+    parse_manifest,
+    parse_result,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -80,9 +82,6 @@ def _filter_session_configuration(
 class SeaDatabricksClient(DatabricksClient):
     """
     Statement Execution API (SEA) implementation of the DatabricksClient interface.
-
-    This implementation provides session management functionality for SEA,
-    while other operations raise NotImplementedError.
     """
 
     # SEA API paths
@@ -92,8 +91,6 @@ class SeaDatabricksClient(DatabricksClient):
     STATEMENT_PATH = BASE_PATH + "statements"
     STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}"
     CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
-    CHUNKS_PATH_WITH_ID = STATEMENT_PATH + "/{}/result/chunks"
-    CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}"
 
     def __init__(
         self,
@@ -126,7 +123,6 @@ def __init__(
         )
 
         self._max_download_threads = kwargs.get("max_download_threads", 10)
-        self.ssl_options = ssl_options
 
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
@@ -283,19 +279,6 @@ def get_default_session_configuration_value(name: str) -> Optional[str]:
         """
         return ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.get(name.upper())
 
-    @staticmethod
-    def is_session_configuration_parameter_supported(name: str) -> bool:
-        """
-        Check if a session configuration parameter is supported.
-
-        Args:
-            name: The name of the session configuration parameter
-
-        Returns:
-            True if the parameter is supported, False otherwise
-        """
-        return name.upper() in ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP
-
     @staticmethod
     def get_allowed_session_configurations() -> List[str]:
         """
@@ -343,92 +326,27 @@ def _results_message_to_execute_response(self, sea_response, command_id):
             tuple: (ExecuteResponse, ResultData, ResultManifest) - The normalized execute response,
                   result data object, and manifest object
         """
-        # Extract status
-        status_data = sea_response.get("status", {})
-        state = CommandState.from_sea_state(status_data.get("state", ""))
-
-        # Extract description from manifest
-        description = None
-        manifest_data = sea_response.get("manifest", {})
-        schema_data = manifest_data.get("schema", {})
-        columns_data = schema_data.get("columns", [])
-
-        if columns_data:
-            columns = []
-            for col_data in columns_data:
-                if not isinstance(col_data, dict):
-                    continue
-
-                # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
-                columns.append(
-                    (
-                        col_data.get("name", ""),  # name
-                        col_data.get("type_name", ""),  # type_code
-                        None,  # display_size (not provided by SEA)
-                        None,  # internal_size (not provided by SEA)
-                        col_data.get("precision"),  # precision
-                        col_data.get("scale"),  # scale
-                        col_data.get("nullable", True),  # null_ok
-                    )
-                )
-            description = columns if columns else None
 
-        # Check for compression
-        lz4_compressed = manifest_data.get("result_compression") == "LZ4_FRAME"
-
-        # Initialize result_data_obj and manifest_obj
-        result_data_obj = None
-        manifest_obj = None
-
-        result_data = sea_response.get("result", {})
-        if result_data:
-            # Convert external links
-            external_links = None
-            if "external_links" in result_data:
-                external_links = []
-                for link_data in result_data["external_links"]:
-                    external_links.append(
-                        ExternalLink(
-                            external_link=link_data.get("external_link", ""),
-                            expiration=link_data.get("expiration", ""),
-                            chunk_index=link_data.get("chunk_index", 0),
-                            byte_count=link_data.get("byte_count", 0),
-                            row_count=link_data.get("row_count", 0),
-                            row_offset=link_data.get("row_offset", 0),
-                            next_chunk_index=link_data.get("next_chunk_index"),
-                            next_chunk_internal_link=link_data.get(
-                                "next_chunk_internal_link"
-                            ),
-                            http_headers=link_data.get("http_headers", {}),
-                        )
-                    )
+        # Parse the response
+        status = parse_status(sea_response)
+        manifest_obj = parse_manifest(sea_response)
+        result_data_obj = parse_result(sea_response)
 
-            # Create the result data object
-            result_data_obj = ResultData(
-                data=result_data.get("data_array"), external_links=external_links
-            )
+        # Extract description from manifest schema
+        description = self._extract_description_from_manifest(manifest_obj)
 
-        # Create the manifest object
-        manifest_obj = ResultManifest(
-            format=manifest_data.get("format", ""),
-            schema=manifest_data.get("schema", {}),
-            total_row_count=manifest_data.get("total_row_count", 0),
-            total_byte_count=manifest_data.get("total_byte_count", 0),
-            total_chunk_count=manifest_data.get("total_chunk_count", 0),
-            truncated=manifest_data.get("truncated", False),
-            chunks=manifest_data.get("chunks"),
-            result_compression=manifest_data.get("result_compression"),
-        )
+        # Check for compression
+        lz4_compressed = manifest_obj.result_compression == "LZ4_FRAME"
 
         execute_response = ExecuteResponse(
             command_id=command_id,
-            status=state,
+            status=status.state,
             description=description,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=False,
             arrow_schema_bytes=None,
-            result_format=manifest_data.get("format"),
+            result_format=manifest_obj.format,
         )
 
         return execute_response, result_data_obj, manifest_obj
@@ -464,6 +382,7 @@ def execute_command(
         Returns:
             ResultSet: A SeaResultSet instance for the executed command
         """
+
         if session_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA session ID")
 
@@ -481,17 +400,25 @@ def execute_command(
                     )
                 )
 
-        format = "ARROW_STREAM" if use_cloud_fetch else "JSON_ARRAY"
-        disposition = "EXTERNAL_LINKS" if use_cloud_fetch else "INLINE"
-        result_compression = "LZ4_FRAME" if lz4_compression else None
+        format = (
+            ResultFormat.ARROW_STREAM if use_cloud_fetch else ResultFormat.JSON_ARRAY
+        ).value
+        disposition = (
+            ResultDisposition.EXTERNAL_LINKS
+            if use_cloud_fetch
+            else ResultDisposition.INLINE
+        ).value
+        result_compression = (
+            ResultCompression.LZ4_FRAME if lz4_compression else ResultCompression.NONE
+        ).value
 
         request = ExecuteStatementRequest(
             warehouse_id=self.warehouse_id,
             session_id=sea_session_id,
             statement=operation,
             disposition=disposition,
             format=format,
-            wait_timeout="0s" if async_op else "10s",
+            wait_timeout=(WaitTimeout.ASYNC if async_op else WaitTimeout.SYNC).value,
             on_wait_timeout="CONTINUE",
             row_limit=max_rows,
             parameters=sea_parameters if sea_parameters else None,
@@ -517,12 +444,11 @@ def execute_command(
         # Store the command ID in the cursor
         cursor.active_command_id = command_id
 
-        # If async operation, return None and let the client poll for results
+        # If async operation, return and let the client poll for results
         if async_op:
             return None
 
         # For synchronous operation, wait for the statement to complete
-        # Poll until the statement is done
         status = response.status
         state = status.state
 
@@ -552,6 +478,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         Raises:
             ValueError: If the command ID is invalid
         """
+
         if command_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA command ID")
 
@@ -574,6 +501,7 @@ def close_command(self, command_id: CommandId) -> None:
         Raises:
             ValueError: If the command ID is invalid
         """
+
         if command_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA command ID")
 
@@ -599,6 +527,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         Raises:
             ValueError: If the command ID is invalid
         """
+
         if command_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA command ID")
 
@@ -633,6 +562,7 @@ def get_execution_result(
         Raises:
             ValueError: If the command ID is invalid
         """
+
         if command_id.backend_type != BackendType.SEA:
             raise ValueError("Not a valid SEA command ID")
 

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

@@ -42,12 +42,29 @@ class ExternalLink:
     http_headers: Optional[Dict[str, str]] = None
 
 
+@dataclass
+class ChunkInfo:
+    """Information about a chunk in the result set."""
+
+    chunk_index: int
+    byte_count: int
+    row_offset: int
+    row_count: int
+
+
 @dataclass
 class ResultData:
     """Result data from a statement execution."""
 
     data: Optional[List[List[Any]]] = None
     external_links: Optional[List[ExternalLink]] = None
+    byte_count: Optional[int] = None
+    chunk_index: Optional[int] = None
+    next_chunk_index: Optional[int] = None
+    next_chunk_internal_link: Optional[str] = None
+    row_count: Optional[int] = None
+    row_offset: Optional[int] = None
+    attachment: Optional[bytes] = None
 
 
 @dataclass
@@ -73,5 +90,6 @@ class ResultManifest:
     total_byte_count: int
     total_chunk_count: int
     truncated: bool = False
-    chunks: Optional[List[Dict[str, Any]]] = None
+    chunks: Optional[List[ChunkInfo]] = None
     result_compression: Optional[str] = None
+    is_volume_operation: Optional[bool] = None

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

@@ -10,7 +10,7 @@
 
 @dataclass
 class StatementParameter:
-    """Parameter for a SQL statement."""
+    """Representation of a parameter for a SQL statement."""
 
     name: str
     value: Optional[str] = None
@@ -19,7 +19,7 @@ class StatementParameter:
 
 @dataclass
 class ExecuteStatementRequest:
-    """Request to execute a SQL statement."""
+    """Representation of a request to execute a SQL statement."""
 
     session_id: str
     statement: str
@@ -65,7 +65,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class GetStatementRequest:
-    """Request to get information about a statement."""
+    """Representation of a request to get information about a statement."""
 
     statement_id: str
 
@@ -76,7 +76,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CancelStatementRequest:
-    """Request to cancel a statement."""
+    """Representation of a request to cancel a statement."""
 
     statement_id: str
 
@@ -87,7 +87,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CloseStatementRequest:
-    """Request to close a statement."""
+    """Representation of a request to close a statement."""
 
     statement_id: str
 
@@ -98,7 +98,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CreateSessionRequest:
-    """Request to create a new session."""
+    """Representation of a request to create a new session."""
 
     warehouse_id: str
     session_confs: Optional[Dict[str, str]] = None
@@ -123,7 +123,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class DeleteSessionRequest:
-    """Request to delete a session."""
+    """Representation of a request to delete a session."""
 
     warehouse_id: str
     session_id: str