use normalised CommandState type in ExecuteResponse

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

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -163,7 +163,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             command_id: The command identifier to check
 
         Returns:
-            ttypes.TOperationState: The current state of the command
+            CommandState: The current state of the command
 
         Raises:
             ValueError: If the command ID is invalid

src/databricks/sql/backend/thrift_backend.py

@@ -86,8 +86,8 @@
 
 
 class ThriftDatabricksClient(DatabricksClient):
-    CLOSED_OP_STATE = ttypes.TOperationState.CLOSED_STATE
-    ERROR_OP_STATE = ttypes.TOperationState.ERROR_STATE
+    CLOSED_OP_STATE = CommandState.CLOSED
+    ERROR_OP_STATE = CommandState.FAILED
 
     _retry_delay_min: float
     _retry_delay_max: float
@@ -351,6 +351,7 @@ 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.
@@ -798,7 +799,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
 
         return ExecuteResponse(
             arrow_queue=arrow_queue_opt,
-            status=operation_state,
+            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,
@@ -863,7 +864,7 @@ def get_execution_result(
 
         execute_response = ExecuteResponse(
             arrow_queue=queue,
-            status=resp.status,
+            status=CommandState.from_thrift_state(resp.status),
             has_been_closed_server_side=False,
             has_more_rows=has_more_rows,
             lz4_compressed=lz4_compressed,

src/databricks/sql/result_set.py

@@ -155,15 +155,12 @@ def __init__(
             arraysize: Default number of rows to fetch
             use_cloud_fetch: Whether to use cloud fetch for retrieving results
         """
-        command_id = execute_response.command_id
-        op_state = CommandState.from_thrift_state(execute_response.status)
-        has_been_closed_server_side = execute_response.has_been_closed_server_side
         super().__init__(
             connection,
             thrift_client,
-            command_id,
-            op_state,
-            has_been_closed_server_side,
+            execute_response.command_id,
+            execute_response.status,
+            execute_response.has_been_closed_server_side,
             arraysize,
             buffer_size_bytes,
         )

tests/e2e/test_driver.py

@@ -827,10 +827,7 @@ def test_close_connection_closes_cursors(self):
                 getProgressUpdate=False,
             )
             op_status_at_server = ars.backend._client.GetOperationStatus(status_request)
-            assert (
-                op_status_at_server.operationState
-                != ttypes.TOperationState.CLOSED_STATE
-            )
+            assert op_status_at_server.operationState != CommandState.CLOSED
 
             conn.close()
 

tests/unit/test_client.py

@@ -96,9 +96,7 @@ def test_closing_connection_closes_commands(self, mock_thrift_client_class):
 
                 mock_execute_response.command_id = Mock(spec=CommandId)
                 mock_execute_response.status = (
-                    TOperationState.FINISHED_STATE
-                    if not closed
-                    else TOperationState.CLOSED_STATE
+                    CommandState.SUCCEEDED if not closed else CommandState.CLOSED
                 )
                 mock_execute_response.has_been_closed_server_side = closed
                 mock_execute_response.is_staging_operation = False

tests/unit/test_thrift_backend.py

@@ -19,7 +19,7 @@
 from databricks.sql.auth.authenticators import AuthProvider
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.result_set import ResultSet, ThriftResultSet
-from databricks.sql.backend.types import CommandId, SessionId, BackendType
+from databricks.sql.backend.types import CommandId, CommandState, SessionId, BackendType
 
 
 def retry_policy_factory():
@@ -883,7 +883,7 @@ def test_handle_execute_response_can_handle_without_direct_results(
                 )
                 self.assertEqual(
                     results_message_response.status,
-                    ttypes.TOperationState.FINISHED_STATE,
+                    CommandState.SUCCEEDED,
                 )
 
     def test_handle_execute_response_can_handle_with_direct_results(self):