Normalise Execution Response (clean backend interfaces)

src/databricks/sql/backend/databricks_client.py

@@ -16,8 +16,6 @@
 
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.types import SessionId, CommandId, CommandState
-from databricks.sql.utils import ExecuteResponse
-from databricks.sql.types import SSLOptions
 
 # Forward reference for type hints
 from typing import TYPE_CHECKING

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, Union, List, Tuple
+from typing import Callable, Dict, Any, Optional, List, Tuple
 from urllib.parse import urljoin
 
 from databricks.sql.auth.authenticators import AuthProvider

src/databricks/sql/backend/thrift_backend.py

@@ -3,23 +3,21 @@
 import logging
 import math
 import time
-import uuid
 import threading
-from typing import List, Optional, Union, Any, TYPE_CHECKING
+from typing import List, Union, Any, TYPE_CHECKING
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
-    from databricks.sql.result_set import ResultSet, ThriftResultSet
 
-from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
     CommandState,
     SessionId,
     CommandId,
-    BackendType,
+    ExecuteResponse,
 )
 from databricks.sql.backend.utils import guid_to_hex_id
 
+
 try:
     import pyarrow
 except ImportError:
@@ -42,7 +40,7 @@
 )
 
 from databricks.sql.utils import (
-    ExecuteResponse,
+    ResultSetQueueFactory,
     _bound,
     RequestErrorInfo,
     NoRetryReason,
@@ -53,6 +51,7 @@
 )
 from databricks.sql.types import SSLOptions
 from databricks.sql.backend.databricks_client import DatabricksClient
+from databricks.sql.result_set import ResultSet, ThriftResultSet
 
 logger = logging.getLogger(__name__)
 
@@ -758,11 +757,13 @@ 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
+
         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
         )
@@ -778,42 +779,28 @@ 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)
 
-        return ExecuteResponse(
-            arrow_queue=arrow_queue_opt,
-            status=CommandState.from_thrift_state(operation_state),
-            has_been_closed_server_side=has_been_closed_server_side,
-            has_more_rows=has_more_rows,
-            lz4_compressed=lz4_compressed,
-            is_staging_operation=is_staging_operation,
+        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 execute_response, has_more_rows
+
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
     ) -> "ResultSet":
-        from databricks.sql.result_set import ThriftResultSet
-
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")
@@ -835,9 +822,6 @@ 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
         )
@@ -852,26 +836,21 @@ def get_execution_result(
         else:
             schema_bytes = None
 
-        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,
-        )
+        lz4_compressed = t_result_set_metadata_resp.lz4Compressed
+        is_staging_operation = t_result_set_metadata_resp.isStagingOperation
+        has_more_rows = resp.hasMoreRows
+
+        status = self.get_query_state(command_id)
 
         execute_response = ExecuteResponse(
-            arrow_queue=queue,
-            status=CommandState.from_thrift_state(resp.status),
+            command_id=command_id,
+            status=status,
+            description=description,
             has_been_closed_server_side=False,
-            has_more_rows=has_more_rows,
             lz4_compressed=lz4_compressed,
             is_staging_operation=is_staging_operation,
-            command_id=command_id,
-            description=description,
             arrow_schema_bytes=schema_bytes,
+            result_format=t_result_set_metadata_resp.resultFormat,
         )
 
         return ThriftResultSet(
@@ -881,6 +860,10 @@ 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,
+            has_more_rows=has_more_rows,
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -947,8 +930,6 @@ def execute_command(
         async_op=False,
         enforce_embedded_schema_correctness=False,
     ) -> Union["ResultSet", None]:
-        from databricks.sql.result_set import ThriftResultSet
-
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -995,7 +976,13 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response = self._handle_execute_response(resp, cursor)
+            execute_response, has_more_rows = 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,
@@ -1004,6 +991,10 @@ def execute_command(
                 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,
+                has_more_rows=has_more_rows,
             )
 
     def get_catalogs(
@@ -1013,8 +1004,6 @@ def get_catalogs(
         max_bytes: int,
         cursor: "Cursor",
     ) -> "ResultSet":
-        from databricks.sql.result_set import ThriftResultSet
-
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1027,7 +1016,11 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = 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,
@@ -1036,6 +1029,10 @@ def get_catalogs(
             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,
+            has_more_rows=has_more_rows,
         )
 
     def get_schemas(
@@ -1047,8 +1044,6 @@ def get_schemas(
         catalog_name=None,
         schema_name=None,
     ) -> "ResultSet":
-        from databricks.sql.result_set import ThriftResultSet
-
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1063,7 +1058,11 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = 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,
@@ -1072,6 +1071,10 @@ def get_schemas(
             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,
+            has_more_rows=has_more_rows,
         )
 
     def get_tables(
@@ -1085,8 +1088,6 @@ def get_tables(
         table_name=None,
         table_types=None,
     ) -> "ResultSet":
-        from databricks.sql.result_set import ThriftResultSet
-
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1103,7 +1104,11 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = 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,
@@ -1112,6 +1117,10 @@ def get_tables(
             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,
+            has_more_rows=has_more_rows,
         )
 
     def get_columns(
@@ -1125,8 +1134,6 @@ def get_columns(
         table_name=None,
         column_name=None,
     ) -> "ResultSet":
-        from databricks.sql.result_set import ThriftResultSet
-
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1143,7 +1150,11 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, has_more_rows = 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,
@@ -1152,6 +1163,10 @@ def get_columns(
             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,
+            has_more_rows=has_more_rows,
         )
 
     def _handle_execute_response(self, resp, cursor):