correct tests and merge artifacts

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -108,7 +108,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         pass
 
     @abstractmethod
-    def close_command(self, command_id: CommandId) -> ttypes.TStatus:
+    def close_command(self, command_id: CommandId) -> None:
         """
         Closes a command and releases associated resources.
 
@@ -118,17 +118,14 @@ def close_command(self, command_id: CommandId) -> ttypes.TStatus:
         Args:
             command_id: The command identifier to close
 
-        Returns:
-            ttypes.TStatus: The status of the close operation
-
         Raises:
             ValueError: If the command ID is invalid
             OperationalError: If there's an error closing the command
         """
         pass
 
     @abstractmethod
-    def get_query_state(self, command_id: CommandId) -> ttypes.TOperationState:
+    def get_query_state(self, command_id: CommandId) -> CommandState:
         """
         Gets the current state of a query or command.
 
@@ -204,7 +201,7 @@ def get_columns(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         column_name: Optional[str] = None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of columns, optionally filtered by catalog, schema, table, and column name patterns.
 

src/databricks/sql/backend/thrift_backend.py

@@ -12,6 +12,7 @@
 
 from databricks.sql.thrift_api.TCLIService.ttypes import TOperationState
 from databricks.sql.backend.types import (
+    CommandState,
     SessionId,
     CommandId,
     BackendType,

src/databricks/sql/backend/types.py

@@ -1,12 +1,46 @@
 from enum import Enum
-from typing import Dict, Optional, Any, Union
+from typing import Dict, Optional, Any
 import logging
 
 from databricks.sql.backend.utils import guid_to_hex_id
+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":
+        if state in (
+            ttypes.TOperationState.INITIALIZED_STATE,
+            ttypes.TOperationState.PENDING_STATE,
+        ):
+            return cls.PENDING
+        elif state == ttypes.TOperationState.RUNNING_STATE:
+            return cls.RUNNING
+        elif state == ttypes.TOperationState.FINISHED_STATE:
+            return cls.SUCCEEDED
+        elif state in (
+            ttypes.TOperationState.ERROR_STATE,
+            ttypes.TOperationState.TIMEDOUT_STATE,
+            ttypes.TOperationState.UKNOWN_STATE,
+        ):
+            return cls.FAILED
+        elif state == ttypes.TOperationState.CLOSED_STATE:
+            return cls.CLOSED
+        elif state == ttypes.TOperationState.CANCELED_STATE:
+            return cls.CANCELLED
+        else:
+            raise ValueError(f"Unknown command state: {state}")
+
+
 class BackendType(Enum):
     """
     Enum representing the type of backend

tests/unit/test_client.py

@@ -133,7 +133,6 @@ def test_closing_result_set_with_closed_connection_soft_closes_commands(self):
 
         result_set = ThriftResultSet(
             connection=mock_connection,
-            backend=mock_backend,
             execute_response=Mock(),
             thrift_client=mock_backend,
         )