Merge branch 'sea-migration' into sea-http-client

Author: varun-edachali-dbx

examples/experimental/tests/test_sea_async_query.py

@@ -45,6 +45,7 @@ def test_sea_async_query_with_cloud_fetch():
             use_sea=True,
             user_agent_entry="SEA-Test-Client",
             use_cloud_fetch=True,
+            enable_query_result_lz4_compression=False,
         )
 
         logger.info(

examples/experimental/tests/test_sea_sync_query.py

@@ -43,6 +43,7 @@ def test_sea_sync_query_with_cloud_fetch():
             use_sea=True,
             user_agent_entry="SEA-Test-Client",
             use_cloud_fetch=True,
+            enable_query_result_lz4_compression=False,
         )
 
         logger.info(

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

@@ -18,7 +18,8 @@
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
-    from databricks.sql.backend.sea.result_set import SeaResultSet
+
+from databricks.sql.backend.sea.result_set import SeaResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (
@@ -130,6 +131,8 @@ def __init__(
             "_use_arrow_native_complex_types", True
         )
 
+        self.use_hybrid_disposition = kwargs.get("use_hybrid_disposition", True)
+
         # Extract warehouse ID from http_path
         self.warehouse_id = self._extract_warehouse_id(http_path)
 
@@ -330,7 +333,7 @@ def _extract_description_from_manifest(
         return columns
 
     def _results_message_to_execute_response(
-        self, response: GetStatementResponse
+        self, response: Union[ExecuteStatementResponse, GetStatementResponse]
     ) -> ExecuteResponse:
         """
         Convert a SEA response to an ExecuteResponse and extract result data.
@@ -364,6 +367,27 @@ def _results_message_to_execute_response(
 
         return execute_response
 
+    def _response_to_result_set(
+        self,
+        response: Union[ExecuteStatementResponse, GetStatementResponse],
+        cursor: Cursor,
+    ) -> SeaResultSet:
+        """
+        Convert a SEA response to a SeaResultSet.
+        """
+
+        execute_response = self._results_message_to_execute_response(response)
+
+        return SeaResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            sea_client=self,
+            result_data=response.result,
+            manifest=response.manifest,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
+        )
+
     def _check_command_not_in_failed_or_closed_state(
         self, state: CommandState, command_id: CommandId
     ) -> None:
@@ -384,21 +408,24 @@ def _check_command_not_in_failed_or_closed_state(
 
     def _wait_until_command_done(
         self, response: ExecuteStatementResponse
-    ) -> CommandState:
+    ) -> Union[ExecuteStatementResponse, GetStatementResponse]:
         """
         Wait until a command is done.
         """
 
-        state = response.status.state
-        command_id = CommandId.from_sea_statement_id(response.statement_id)
+        final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response
+
+        state = final_response.status.state
+        command_id = CommandId.from_sea_statement_id(final_response.statement_id)
 
         while state in [CommandState.PENDING, CommandState.RUNNING]:
             time.sleep(self.POLL_INTERVAL_SECONDS)
-            state = self.get_query_state(command_id)
+            final_response = self._poll_query(command_id)
+            state = final_response.status.state
 
         self._check_command_not_in_failed_or_closed_state(state, command_id)
 
-        return state
+        return final_response
 
     def execute_command(
         self,
@@ -456,7 +483,11 @@ def execute_command(
             ResultFormat.ARROW_STREAM if use_cloud_fetch else ResultFormat.JSON_ARRAY
         ).value
         disposition = (
-            ResultDisposition.EXTERNAL_LINKS
+            (
+                ResultDisposition.HYBRID
+                if self.use_hybrid_disposition
+                else ResultDisposition.EXTERNAL_LINKS
+            )
             if use_cloud_fetch
             else ResultDisposition.INLINE
         ).value
@@ -500,8 +531,11 @@ def execute_command(
         if async_op:
             return None
 
-        self._wait_until_command_done(response)
-        return self.get_execution_result(command_id, cursor)
+        final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response
+        if response.status.state != CommandState.SUCCEEDED:
+            final_response = self._wait_until_command_done(response)
+
+        return self._response_to_result_set(final_response, cursor)
 
     def cancel_command(self, command_id: CommandId) -> None:
         """
@@ -553,18 +587,9 @@ def close_command(self, command_id: CommandId) -> None:
             data=request.to_dict(),
         )
 
-    def get_query_state(self, command_id: CommandId) -> CommandState:
+    def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
         """
-        Get the state of a running query.
-
-        Args:
-            command_id: Command identifier
-
-        Returns:
-            CommandState: The current state of the command
-
-        Raises:
-            ValueError: If the command ID is invalid
+        Poll for the current command info.
         """
 
         if command_id.backend_type != BackendType.SEA:
@@ -580,9 +605,25 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
         )
-
-        # Parse the response
         response = GetStatementResponse.from_dict(response_data)
+
+        return response
+
+    def get_query_state(self, command_id: CommandId) -> CommandState:
+        """
+        Get the state of a running query.
+
+        Args:
+            command_id: Command identifier
+
+        Returns:
+            CommandState: The current state of the command
+
+        Raises:
+            ProgrammingError: If the command ID is invalid
+        """
+
+        response = self._poll_query(command_id)
         return response.status.state
 
     def get_execution_result(
@@ -604,40 +645,12 @@ def get_execution_result(
             ValueError: If the command ID is invalid
         """
 
-        if command_id.backend_type != BackendType.SEA:
-            raise ValueError("Not a valid SEA command ID")
-
-        sea_statement_id = command_id.to_sea_statement_id()
-        if sea_statement_id is None:
-            raise ValueError("Not a valid SEA command ID")
-
-        # Create the request model
-        request = GetStatementRequest(statement_id=sea_statement_id)
-
-        # Get the statement result
-        response_data = self._http_client._make_request(
-            method="GET",
-            path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
-            data=request.to_dict(),
-        )
-        response = GetStatementResponse.from_dict(response_data)
-
-        # Create and return a SeaResultSet
-        from databricks.sql.backend.sea.result_set import SeaResultSet
+        response = self._poll_query(command_id)
+        return self._response_to_result_set(response, cursor)
 
-        execute_response = self._results_message_to_execute_response(response)
-
-        return SeaResultSet(
-            connection=cursor.connection,
-            execute_response=execute_response,
-            sea_client=self,
-            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:
+    def get_chunk_links(
+        self, statement_id: str, chunk_index: int
+    ) -> List[ExternalLink]:
         """
         Get links for chunks starting from the specified index.
         Args:
@@ -654,17 +667,7 @@ def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
         response = GetChunksResponse.from_dict(response_data)
 
         links = response.external_links or []
-        link = next((l for l in links if l.chunk_index == chunk_index), None)
-        if not link:
-            raise ServerOperationError(
-                f"No link found for chunk index {chunk_index}",
-                {
-                    "operation-id": statement_id,
-                    "diagnostic-info": None,
-                },
-            )
-
-        return link
+        return links
 
     # == Metadata Operations ==
 

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

@@ -4,6 +4,7 @@
 These models define the structures used in SEA API responses.
 """
 
+import base64
 from typing import Dict, Any, List, Optional
 from dataclasses import dataclass
 
@@ -91,6 +92,11 @@ def _parse_result(data: Dict[str, Any]) -> ResultData:
                 )
             )
 
+    # Handle attachment field - decode from base64 if present
+    attachment = result_data.get("attachment")
+    if attachment is not None:
+        attachment = base64.b64decode(attachment)
+
     return ResultData(
         data=result_data.get("data_array"),
         external_links=external_links,
@@ -100,7 +106,7 @@ def _parse_result(data: Dict[str, Any]) -> ResultData:
         next_chunk_internal_link=result_data.get("next_chunk_internal_link"),
         row_count=result_data.get("row_count"),
         row_offset=result_data.get("row_offset"),
-        attachment=result_data.get("attachment"),
+        attachment=attachment,
     )
 
 

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

@@ -5,6 +5,8 @@
 
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 
+from databricks.sql.cloudfetch.downloader import ResultSetDownloadHandler
+
 try:
     import pyarrow
 except ImportError:
@@ -23,7 +25,12 @@
 from databricks.sql.exc import ProgrammingError, ServerOperationError
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink
 from databricks.sql.types import SSLOptions
-from databricks.sql.utils import CloudFetchQueue, ResultSetQueue
+from databricks.sql.utils import (
+    ArrowQueue,
+    CloudFetchQueue,
+    ResultSetQueue,
+    create_arrow_table_from_arrow_file,
+)
 
 import logging
 
@@ -62,6 +69,18 @@ def build_queue(
             # INLINE disposition with JSON_ARRAY format
             return JsonQueue(result_data.data)
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
+            if result_data.attachment is not None:
+                arrow_file = (
+                    ResultSetDownloadHandler._decompress_data(result_data.attachment)
+                    if lz4_compressed
+                    else result_data.attachment
+                )
+                arrow_table = create_arrow_table_from_arrow_file(
+                    arrow_file, description
+                )
+                logger.debug(f"Created arrow table with {arrow_table.num_rows} rows")
+                return ArrowQueue(arrow_table, manifest.total_row_count)
+
             # EXTERNAL_LINKS disposition
             return SeaCloudFetchQueue(
                 result_data=result_data,
@@ -150,7 +169,11 @@ def __init__(
         )
 
         initial_links = result_data.external_links or []
-        first_link = next((l for l in initial_links if l.chunk_index == 0), None)
+        self._chunk_index_to_link = {link.chunk_index: link for link in initial_links}
+
+        # Track the current chunk we're processing
+        self._current_chunk_index = 0
+        first_link = self._chunk_index_to_link.get(self._current_chunk_index, None)
         if not first_link:
             # possibly an empty response
             return None
@@ -173,21 +196,24 @@ def _convert_to_thrift_link(self, link: ExternalLink) -> TSparkArrowResultLink:
             httpHeaders=link.http_headers or {},
         )
 
-    def _get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
-        """Progress to the next chunk link."""
+    def _get_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
         if chunk_index >= self._total_chunk_count:
             return None
 
-        try:
-            return self._sea_client.get_chunk_link(self._statement_id, chunk_index)
-        except Exception as e:
+        if chunk_index not in self._chunk_index_to_link:
+            links = self._sea_client.get_chunk_links(self._statement_id, chunk_index)
+            self._chunk_index_to_link.update({l.chunk_index: l for l in links})
+
+        link = self._chunk_index_to_link.get(chunk_index, None)
+        if not link:
             raise ServerOperationError(
-                f"Error fetching link for chunk {chunk_index}: {e}",
+                f"Error fetching link for chunk {chunk_index}",
                 {
                     "operation-id": self._statement_id,
                     "diagnostic-info": None,
                 },
             )
+        return link
 
     def _create_table_from_link(
         self, link: ExternalLink

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

@@ -4,7 +4,6 @@
 
 import logging
 
-from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
 from databricks.sql.backend.sea.utils.conversion import SqlTypeConverter
 
@@ -15,6 +14,7 @@
 
 if TYPE_CHECKING:
     from databricks.sql.client import Connection
+    from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.types import Row
 from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.types import ExecuteResponse

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

@@ -28,7 +28,7 @@ class ResultFormat(Enum):
 class ResultDisposition(Enum):
     """Enum for result disposition values."""
 
-    # TODO: add support for hybrid disposition
+    HYBRID = "INLINE_OR_EXTERNAL_LINKS"
     EXTERNAL_LINKS = "EXTERNAL_LINKS"
     INLINE = "INLINE"
 

src/databricks/sql/client.py

@@ -99,6 +99,10 @@ def __init__(
         Connect to a Databricks SQL endpoint or a Databricks cluster.
 
         Parameters:
+            :param use_sea: `bool`, optional (default is False)
+                Use the SEA backend instead of the Thrift backend.
+            :param use_hybrid_disposition: `bool`, optional (default is False)
+                Use the hybrid disposition instead of the inline disposition.
             :param server_hostname: Databricks instance host name.
             :param http_path: Http path either to a DBSQL endpoint (e.g. /sql/1.0/endpoints/1234567890abcdef)
                 or to a DBR interactive cluster (e.g. /sql/protocolv1/o/1234567890123456/1234-123456-slid123)