Revert "Merge branch 'sea-migration' into exec-models-sea"

This reverts commit 8bd12d8, reversing
changes made to 030edf8.

Author: varun-edachali-dbx

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

@@ -98,7 +98,7 @@ def to_dict(self) -> Dict[str, Any]:
 
 @dataclass
 class CreateSessionRequest:
-    """Representation of a request to create a new session."""
+    """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:
-    """Representation of a request to delete a session."""
+    """Request to delete a session."""
 
     warehouse_id: str
     session_id: str

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

@@ -146,7 +146,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
 
 @dataclass
 class CreateSessionResponse:
-    """Representation of the response from creating a new session."""
+    """Response from creating a new session."""
 
     session_id: str
 

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

@@ -1,7 +1,7 @@
 import json
 import logging
 import requests
-from typing import Callable, Dict, Any, Optional, List, Tuple
+from typing import Callable, Dict, Any, Optional, Union, List, Tuple
 from urllib.parse import urljoin
 
 from databricks.sql.auth.authenticators import AuthProvider

src/databricks/sql/backend/thrift_backend.py

@@ -3,21 +3,24 @@
 import logging
 import math
 import time
+import uuid
 import threading
 from typing import List, Union, Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
+from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
     CommandState,
     SessionId,
     CommandId,
+    BackendType,
+    guid_to_hex_id,
     ExecuteResponse,
 )
 from databricks.sql.backend.utils import guid_to_hex_id
 
-
 try:
     import pyarrow
 except ImportError:
@@ -757,13 +760,11 @@ def _results_message_to_execute_response(self, resp, operation_state):
             )
         direct_results = resp.directResults
         has_been_closed_server_side = direct_results and direct_results.closeOperation
-
-        is_direct_results = (
+        has_more_rows = (
             (not direct_results)
             or (not direct_results.resultSet)
             or direct_results.resultSet.hasMoreRows
         )
-
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema
         )
@@ -779,25 +780,43 @@ def _results_message_to_execute_response(self, resp, operation_state):
             schema_bytes = None
 
         lz4_compressed = t_result_set_metadata_resp.lz4Compressed
+        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
+        if direct_results and direct_results.resultSet:
+            assert direct_results.resultSet.results.startRowOffset == 0
+            assert direct_results.resultSetMetadata
+
+            arrow_queue_opt = ResultSetQueueFactory.build_queue(
+                row_set_type=t_result_set_metadata_resp.resultFormat,
+                t_row_set=direct_results.resultSet.results,
+                arrow_schema_bytes=schema_bytes,
+                max_download_threads=self.max_download_threads,
+                lz4_compressed=lz4_compressed,
+                description=description,
+                ssl_options=self._ssl_options,
+            )
+        else:
+            arrow_queue_opt = None
+
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
 
         status = CommandState.from_thrift_state(operation_state)
         if status is None:
             raise ValueError(f"Unknown command state: {operation_state}")
 
-        execute_response = ExecuteResponse(
-            command_id=command_id,
-            status=status,
-            description=description,
-            has_been_closed_server_side=has_been_closed_server_side,
-            lz4_compressed=lz4_compressed,
-            is_staging_operation=t_result_set_metadata_resp.isStagingOperation,
-            arrow_schema_bytes=schema_bytes,
-            result_format=t_result_set_metadata_resp.resultFormat,
+        return (
+            ExecuteResponse(
+                command_id=command_id,
+                status=status,
+                description=description,
+                has_more_rows=has_more_rows,
+                results_queue=arrow_queue_opt,
+                has_been_closed_server_side=has_been_closed_server_side,
+                lz4_compressed=lz4_compressed,
+                is_staging_operation=is_staging_operation,
+            ),
+            schema_bytes,
         )
 
-        return execute_response, is_direct_results
-
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
     ) -> "ResultSet":
@@ -822,6 +841,9 @@ def get_execution_result(
 
         t_result_set_metadata_resp = resp.resultSetMetadata
 
+        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
+        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
+        has_more_rows = resp.hasMoreRows
         description = self._hive_schema_to_description(
             t_result_set_metadata_resp.schema
         )
@@ -836,21 +858,25 @@ def get_execution_result(
         else:
             schema_bytes = None
 
-        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
-        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
-        is_direct_results = resp.hasMoreRows
-
-        status = self.get_query_state(command_id)
+        queue = ResultSetQueueFactory.build_queue(
+            row_set_type=resp.resultSetMetadata.resultFormat,
+            t_row_set=resp.results,
+            arrow_schema_bytes=schema_bytes,
+            max_download_threads=self.max_download_threads,
+            lz4_compressed=lz4_compressed,
+            description=description,
+            ssl_options=self._ssl_options,
+        )
 
         execute_response = ExecuteResponse(
             command_id=command_id,
             status=status,
             description=description,
+            has_more_rows=has_more_rows,
+            results_queue=queue,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=is_staging_operation,
-            arrow_schema_bytes=schema_bytes,
-            result_format=t_result_set_metadata_resp.resultFormat,
         )
 
         return ThriftResultSet(
@@ -860,10 +886,7 @@ def get_execution_result(
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            t_row_set=resp.results,
-            max_download_threads=self.max_download_threads,
-            ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            arrow_schema_bytes=schema_bytes,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -976,25 +999,18 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response, is_direct_results = self._handle_execute_response(
+            execute_response, arrow_schema_bytes = self._handle_execute_response(
                 resp, cursor
             )
 
-            t_row_set = None
-            if resp.directResults and resp.directResults.resultSet:
-                t_row_set = resp.directResults.resultSet.results
-
             return ThriftResultSet(
                 connection=cursor.connection,
                 execute_response=execute_response,
                 thrift_client=self,
                 buffer_size_bytes=max_bytes,
                 arraysize=max_rows,
                 use_cloud_fetch=use_cloud_fetch,
-                t_row_set=t_row_set,
-                max_download_threads=self.max_download_threads,
-                ssl_options=self._ssl_options,
-                is_direct_results=is_direct_results,
+                arrow_schema_bytes=arrow_schema_bytes,
             )
 
     def get_catalogs(
@@ -1016,25 +1032,18 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
             resp, cursor
         )
 
-        t_row_set = None
-        if resp.directResults and resp.directResults.resultSet:
-            t_row_set = resp.directResults.resultSet.results
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            t_row_set=t_row_set,
-            max_download_threads=self.max_download_threads,
-            ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def get_schemas(
@@ -1060,25 +1069,18 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
             resp, cursor
         )
 
-        t_row_set = None
-        if resp.directResults and resp.directResults.resultSet:
-            t_row_set = resp.directResults.resultSet.results
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            t_row_set=t_row_set,
-            max_download_threads=self.max_download_threads,
-            ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def get_tables(
@@ -1108,25 +1110,18 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
             resp, cursor
         )
 
-        t_row_set = None
-        if resp.directResults and resp.directResults.resultSet:
-            t_row_set = resp.directResults.resultSet.results
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            t_row_set=t_row_set,
-            max_download_threads=self.max_download_threads,
-            ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def get_columns(
@@ -1156,25 +1151,18 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response, is_direct_results = self._handle_execute_response(
+        execute_response, arrow_schema_bytes = self._handle_execute_response(
             resp, cursor
         )
 
-        t_row_set = None
-        if resp.directResults and resp.directResults.resultSet:
-            t_row_set = resp.directResults.resultSet.results
-
         return ThriftResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             thrift_client=self,
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
-            t_row_set=t_row_set,
-            max_download_threads=self.max_download_threads,
-            ssl_options=self._ssl_options,
-            is_direct_results=is_direct_results,
+            arrow_schema_bytes=arrow_schema_bytes,
         )
 
     def _handle_execute_response(self, resp, cursor):

src/databricks/sql/backend/types.py

@@ -423,9 +423,11 @@ class ExecuteResponse:
 
     command_id: CommandId
     status: CommandState
-    description: Optional[List[Tuple]] = None
+    description: Optional[
+        List[Tuple[str, str, None, None, Optional[int], Optional[int], bool]]
+    ] = None
+    has_more_rows: bool = False
+    results_queue: Optional[Any] = None
     has_been_closed_server_side: bool = False
     lz4_compressed: bool = True
     is_staging_operation: bool = False
-    arrow_schema_bytes: Optional[bytes] = None
-    result_format: Optional[Any] = None