account for new exec resp

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/filters.py

@@ -18,7 +18,9 @@
 )
 
 if TYPE_CHECKING:
-    from databricks.sql.result_set import ResultSet, SeaResultSet
+    from databricks.sql.result_set import ResultSet
+
+from databricks.sql.result_set import SeaResultSet
 
 logger = logging.getLogger(__name__)
 

src/databricks/sql/backend/thrift_backend.py

@@ -5,20 +5,20 @@
 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,
+    guid_to_hex_id,
+    ExecuteResponse,
 )
-from databricks.sql.backend.utils import guid_to_hex_id
 
 try:
     import pyarrow
@@ -42,7 +42,7 @@
 )
 
 from databricks.sql.utils import (
-    ExecuteResponse,
+    ResultSetQueueFactory,
     _bound,
     RequestErrorInfo,
     NoRetryReason,
@@ -53,6 +53,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__)
 
@@ -351,7 +352,6 @@ def make_request(self, method, request, retryable=True):
         Will stop retry attempts if total elapsed time + next retry delay would exceed
         _retry_stop_after_attempts_duration.
         """
-
         # basic strategy: build range iterator rep'ing number of available
         # retries. bounds can be computed from there. iterate over it with
         # retries until success or final failure achieved.
@@ -797,23 +797,24 @@ def _results_message_to_execute_response(self, resp, operation_state):
 
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
 
+        status = CommandState.from_thrift_state(operation_state)
+        if status is None:
+            raise ValueError(f"Invalid operation state: {operation_state}")
+
         return ExecuteResponse(
-            arrow_queue=arrow_queue_opt,
-            status=CommandState.from_thrift_state(operation_state),
-            has_been_closed_server_side=has_been_closed_server_side,
+            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,
-            command_id=command_id,
-            description=description,
-            arrow_schema_bytes=schema_bytes,
-        )
+        ), schema_bytes
 
     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")
@@ -863,15 +864,14 @@ def get_execution_result(
         )
 
         execute_response = ExecuteResponse(
-            arrow_queue=queue,
-            status=CommandState.from_thrift_state(resp.status),
-            has_been_closed_server_side=False,
+            command_id=command_id,
+            status=resp.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,
-            command_id=command_id,
-            description=description,
-            arrow_schema_bytes=schema_bytes,
         )
 
         return ThriftResultSet(
@@ -881,6 +881,7 @@ def get_execution_result(
             buffer_size_bytes=cursor.buffer_size_bytes,
             arraysize=cursor.arraysize,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            arrow_schema_bytes=schema_bytes
         )
 
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
@@ -909,10 +910,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         poll_resp = self._poll_for_status(thrift_handle)
         operation_state = poll_resp.operationState
         self._check_command_not_in_error_or_closed_state(thrift_handle, poll_resp)
-        state = CommandState.from_thrift_state(operation_state)
-        if state is None:
-            raise ValueError(f"Unknown command state: {operation_state}")
-        return state
+        return CommandState.from_thrift_state(operation_state)
 
     @staticmethod
     def _check_direct_results_for_error(t_spark_direct_results):
@@ -947,8 +945,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 +991,7 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            execute_response = self._handle_execute_response(resp, cursor)
+            execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
 
             return ThriftResultSet(
                 connection=cursor.connection,
@@ -1004,6 +1000,7 @@ def execute_command(
                 buffer_size_bytes=max_bytes,
                 arraysize=max_rows,
                 use_cloud_fetch=use_cloud_fetch,
+                arrow_schema_bytes=arrow_schema_bytes
             )
 
     def get_catalogs(
@@ -1013,8 +1010,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 +1022,7 @@ def get_catalogs(
         )
         resp = self.make_request(self._client.GetCatalogs, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1036,6 +1031,7 @@ def get_catalogs(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            arrow_schema_bytes=arrow_schema_bytes
         )
 
     def get_schemas(
@@ -1047,8 +1043,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 +1057,7 @@ def get_schemas(
         )
         resp = self.make_request(self._client.GetSchemas, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1072,6 +1066,7 @@ def get_schemas(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            arrow_schema_bytes=arrow_schema_bytes
         )
 
     def get_tables(
@@ -1085,8 +1080,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 +1096,7 @@ def get_tables(
         )
         resp = self.make_request(self._client.GetTables, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1112,6 +1105,7 @@ def get_tables(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            arrow_schema_bytes=arrow_schema_bytes
         )
 
     def get_columns(
@@ -1125,8 +1119,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 +1135,7 @@ def get_columns(
         )
         resp = self.make_request(self._client.GetColumns, req)
 
-        execute_response = self._handle_execute_response(resp, cursor)
+        execute_response, arrow_schema_bytes = self._handle_execute_response(resp, cursor)
 
         return ThriftResultSet(
             connection=cursor.connection,
@@ -1152,6 +1144,7 @@ def get_columns(
             buffer_size_bytes=max_bytes,
             arraysize=max_rows,
             use_cloud_fetch=cursor.connection.use_cloud_fetch,
+            arrow_schema_bytes=arrow_schema_bytes
         )
 
     def _handle_execute_response(self, resp, cursor):
@@ -1165,11 +1158,11 @@ def _handle_execute_response(self, resp, cursor):
             resp.directResults and resp.directResults.operationStatus,
         )
 
-        execute_response = self._results_message_to_execute_response(
+        execute_response, arrow_schema_bytes = self._results_message_to_execute_response(
             resp, final_operation_state
         )
-        execute_response = execute_response._replace(command_id=command_id)
-        return execute_response
+        execute_response.command_id = command_id
+        return execute_response, arrow_schema_bytes
 
     def _handle_execute_response_async(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
@@ -1230,7 +1223,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")
 
-        logger.debug("Cancelling command {}".format(guid_to_hex_id(command_id.guid)))
+        logger.debug("Cancelling command {}".format(command_id.guid))
         req = ttypes.TCancelOperationReq(thrift_handle)
         self.make_request(self._client.CancelOperation, req)
 

tests/unit/test_fetches.py

@@ -43,14 +43,13 @@ def make_dummy_result_set_from_initial_results(initial_results):
         rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
+                command_id=None,
                 status=None,
                 has_been_closed_server_side=True,
                 has_more_rows=False,
                 description=Mock(),
                 lz4_compressed=Mock(),
-                command_id=None,
                 results_queue=arrow_queue,
-                arrow_schema_bytes=schema.serialize().to_pybytes(),
                 is_staging_operation=False,
             ),
             thrift_client=None,
@@ -89,6 +88,7 @@ def fetch_results(
         rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
+                command_id=None,
                 status=None,
                 has_been_closed_server_side=False,
                 has_more_rows=True,
@@ -97,9 +97,7 @@ def fetch_results(
                     for col_id in range(num_cols)
                 ],
                 lz4_compressed=Mock(),
-                command_id=None,
                 results_queue=None,
-                arrow_schema_bytes=None,
                 is_staging_operation=False,
             ),
             thrift_client=mock_thrift_backend,