Implement ResultSet Abstraction (backend interfaces for fetch phase)

src/databricks/sql/backend/databricks_client.py

@@ -15,10 +15,16 @@
     from databricks.sql.client import Cursor
 
 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
 
+# Forward reference for type hints
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from databricks.sql.result_set import ResultSet
+
 
 class DatabricksClient(ABC):
     # == Connection and Session Management ==
@@ -81,7 +87,7 @@ def execute_command(
         parameters: List[ttypes.TSparkParameter],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
-    ) -> Optional[ExecuteResponse]:
+    ) -> Union["ResultSet", None]:
         """
         Executes a SQL command or query within the specified session.
 
@@ -130,7 +136,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.
 
@@ -140,17 +146,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.
 
@@ -175,7 +178,7 @@ def get_execution_result(
         self,
         command_id: CommandId,
         cursor: "Cursor",
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves the results of a previously executed command.
 
@@ -203,7 +206,7 @@ def get_catalogs(
         max_rows: int,
         max_bytes: int,
         cursor: "Cursor",
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of available catalogs.
 
@@ -234,7 +237,7 @@ def get_schemas(
         cursor: "Cursor",
         catalog_name: Optional[str] = None,
         schema_name: Optional[str] = None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of schemas, optionally filtered by catalog and schema name patterns.
 
@@ -269,7 +272,7 @@ def get_tables(
         schema_name: Optional[str] = None,
         table_name: Optional[str] = None,
         table_types: Optional[List[str]] = None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         """
         Retrieves a list of tables, optionally filtered by catalog, schema, table name, and table types.
 
@@ -306,7 +309,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,
@@ -51,6 +52,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__)
 
@@ -808,7 +810,7 @@ def _results_message_to_execute_response(self, resp, operation_state):
 
     def get_execution_result(
         self, command_id: CommandId, cursor: "Cursor"
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")
@@ -857,7 +859,7 @@ def get_execution_result(
             ssl_options=self._ssl_options,
         )
 
-        return ExecuteResponse(
+        execute_response = ExecuteResponse(
             arrow_queue=queue,
             status=resp.status,
             has_been_closed_server_side=False,
@@ -869,6 +871,15 @@ def get_execution_result(
             arrow_schema_bytes=schema_bytes,
         )
 
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
+        )
+
     def _wait_until_command_done(self, op_handle, initial_operation_status_resp):
         if initial_operation_status_resp:
             self._check_command_not_in_error_or_closed_state(
@@ -887,15 +898,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):
@@ -929,7 +940,7 @@ def execute_command(
         parameters=[],
         async_op=False,
         enforce_embedded_schema_correctness=False,
-    ) -> Optional[ExecuteResponse]:
+    ) -> Union["ResultSet", None]:
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -976,15 +987,24 @@ def execute_command(
             self._handle_execute_response_async(resp, cursor)
             return None
         else:
-            return self._handle_execute_response(resp, cursor)
+            execute_response = self._handle_execute_response(resp, cursor)
+
+            return ThriftResultSet(
+                connection=cursor.connection,
+                execute_response=execute_response,
+                thrift_client=self,
+                buffer_size_bytes=max_bytes,
+                arraysize=max_rows,
+                use_cloud_fetch=use_cloud_fetch,
+            )
 
     def get_catalogs(
         self,
         session_id: SessionId,
         max_rows: int,
         max_bytes: int,
         cursor: "Cursor",
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -996,7 +1016,17 @@ def get_catalogs(
             ),
         )
         resp = self.make_request(self._client.GetCatalogs, req)
-        return self._handle_execute_response(resp, cursor)
+
+        execute_response = self._handle_execute_response(resp, cursor)
+
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
+        )
 
     def get_schemas(
         self,
@@ -1006,7 +1036,7 @@ def get_schemas(
         cursor: "Cursor",
         catalog_name=None,
         schema_name=None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1020,7 +1050,17 @@ def get_schemas(
             schemaName=schema_name,
         )
         resp = self.make_request(self._client.GetSchemas, req)
-        return self._handle_execute_response(resp, cursor)
+
+        execute_response = self._handle_execute_response(resp, cursor)
+
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
+        )
 
     def get_tables(
         self,
@@ -1032,7 +1072,7 @@ def get_tables(
         schema_name=None,
         table_name=None,
         table_types=None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1048,7 +1088,17 @@ def get_tables(
             tableTypes=table_types,
         )
         resp = self.make_request(self._client.GetTables, req)
-        return self._handle_execute_response(resp, cursor)
+
+        execute_response = self._handle_execute_response(resp, cursor)
+
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
+        )
 
     def get_columns(
         self,
@@ -1060,7 +1110,7 @@ def get_columns(
         schema_name=None,
         table_name=None,
         column_name=None,
-    ) -> ExecuteResponse:
+    ) -> "ResultSet":
         thrift_handle = session_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift session ID")
@@ -1076,7 +1126,17 @@ def get_columns(
             columnName=column_name,
         )
         resp = self.make_request(self._client.GetColumns, req)
-        return self._handle_execute_response(resp, cursor)
+
+        execute_response = self._handle_execute_response(resp, cursor)
+
+        return ThriftResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            thrift_client=self,
+            buffer_size_bytes=max_bytes,
+            arraysize=max_rows,
+            use_cloud_fetch=cursor.connection.use_cloud_fetch,
+        )
 
     def _handle_execute_response(self, resp, cursor):
         command_id = CommandId.from_thrift_handle(resp.operationHandle)
@@ -1154,7 +1214,7 @@ def cancel_command(self, command_id: CommandId) -> None:
         req = ttypes.TCancelOperationReq(thrift_handle)
         self.make_request(self._client.CancelOperation, req)
 
-    def close_command(self, command_id: CommandId):
+    def close_command(self, command_id: CommandId) -> None:
         thrift_handle = command_id.to_thrift_handle()
         if not thrift_handle:
             raise ValueError("Not a valid Thrift command ID")

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
@@ -211,6 +245,7 @@ def __str__(self) -> str:
         Returns:
             A string representation of the command ID
         """
+
         if self.backend_type == BackendType.SEA:
             return str(self.guid)
         elif self.backend_type == BackendType.THRIFT: