abstract Command State away from Thrift specific types

ref:
https://github.com/databricks/databricks-jdbc/blob/30688befd95060874b2dd00019dc85a449ed435a/src/main/java/com/databricks/jdbc/common/util/DatabricksThriftUtil.java#L85
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -2,7 +2,7 @@
 from typing import Dict, Tuple, List, Optional, Any, Union
 
 from databricks.sql.thrift_api.TCLIService import ttypes
-from databricks.sql.backend.types import SessionId, CommandId
+from databricks.sql.backend.types import SessionId, CommandId, CommandState
 from databricks.sql.utils import ExecuteResponse
 from databricks.sql.types import SSLOptions
 
@@ -54,7 +54,7 @@ def close_command(self, command_id: CommandId) -> ttypes.TStatus:
         pass
 
     @abstractmethod
-    def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
+    def get_query_state(self, command_id: CommandId) -> CommandState:
         pass
 
     @abstractmethod

src/databricks/sql/backend/thrift_backend.py

@@ -9,6 +9,7 @@
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
+    CommandState,
     SessionId,
     CommandId,
     BackendType,
@@ -903,15 +904,15 @@ def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
             self._check_command_not_in_error_or_closed_state(op_handle, poll_resp)
         return operation_state
 
-    def get_query_state(self, command_id: CommandId) -> "TOperationState":
+    def get_query_state(self, command_id: CommandId) -> CommandState:
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")
 
         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)
-        return operation_state
+        return CommandState.from_thrift_state(operation_state)
 
     @staticmethod
     def _check_direct_results_for_error(t_spark_direct_results):

src/databricks/sql/backend/types.py

@@ -3,9 +3,38 @@
 import uuid
 import logging
 
+from databricks.sql.thrift_api.TCLIService import ttypes
+
 logger = logging.getLogger(__name__)
 
 
+class CommandState(Enum):
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    SUCCEEDED = "SUCCEEDED"
+    FAILED = "FAILED"
+    CLOSED = "CLOSED"
+    CANCELLED = "CANCELLED"
+
+    @classmethod
+    def from_thrift_state(cls, state: ttypes.TOperationState) -> "CommandState":
+        match state:
+            case ttypes.TOperationState.INITIALIZED_STATE | ttypes.TOperationState.PENDING_STATE:
+                return cls.PENDING
+            case ttypes.TOperationState.RUNNING_STATE:
+                return cls.RUNNING
+            case ttypes.TOperationState.FINISHED_STATE:
+                return cls.SUCCEEDED
+            case ttypes.TOperationState.ERROR_STATE | ttypes.TOperationState.TIMEDOUT_STATE | ttypes.TOperationState.UKNOWN_STATE:
+                return cls.FAILED
+            case ttypes.TOperationState.CLOSED_STATE:
+                return cls.CLOSED
+            case ttypes.TOperationState.CANCELLED_STATE:
+                return cls.CANCELLED
+            case _:
+                raise ValueError(f"Unknown command state: {state}")
+
+
 def guid_to_hex_id(guid: bytes) -> str:
     """Return a hexadecimal string instead of bytes
 

src/databricks/sql/client.py

@@ -47,7 +47,7 @@
 from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
 from databricks.sql.experimental.oauth_persistence import OAuthPersistence
 from databricks.sql.session import Session
-from databricks.sql.backend.types import CommandId, BackendType, SessionId
+from databricks.sql.backend.types import CommandId, BackendType, CommandState, SessionId
 
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TSparkParameter,
@@ -846,7 +846,7 @@ def execute_async(
 
         return self
 
-    def get_query_state(self) -> "TOperationState":
+    def get_query_state(self) -> CommandState:
         """
         Get the state of the async executing query or basically poll the status of the query
 
@@ -862,11 +862,7 @@ def is_query_pending(self):
         :return:
         """
         operation_state = self.get_query_state()
-
-        return not operation_state or operation_state in [
-            ttypes.TOperationState.RUNNING_STATE,
-            ttypes.TOperationState.PENDING_STATE,
-        ]
+        return operation_state in [CommandState.PENDING, CommandState.RUNNING]
 
     def get_async_execution_result(self):
         """
@@ -882,7 +878,7 @@ def get_async_execution_result(self):
             time.sleep(self.ASYNC_DEFAULT_POLLING_INTERVAL)
 
         operation_state = self.get_query_state()
-        if operation_state == ttypes.TOperationState.FINISHED_STATE:
+        if operation_state == CommandState.SUCCEEDED:
             self.active_result_set = self.backend.get_execution_result(
                 self.active_op_handle, self
             )