move common state to ResultSet aprent

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

Author: varun-edachali-dbx

src/databricks/sql/backend/thrift_backend.py

@@ -9,6 +9,7 @@
 
 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 (
@@ -52,7 +53,6 @@
 )
 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__)
 
@@ -811,6 +811,8 @@ def _results_message_to_execute_response(self, resp, operation_state):
     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")
@@ -906,7 +908,10 @@ 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)
-        return CommandState.from_thrift_state(operation_state)
+        state = CommandState.from_thrift_state(operation_state)
+        if state is None:
+            raise ValueError(f"Unknown command state: {operation_state}")
+        return state
 
     @staticmethod
     def _check_direct_results_for_error(t_spark_direct_results):
@@ -941,6 +946,8 @@ 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")
@@ -1005,6 +1012,8 @@ 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")
@@ -1037,6 +1046,8 @@ 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")
@@ -1073,6 +1084,8 @@ 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")
@@ -1111,6 +1124,8 @@ 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")

src/databricks/sql/backend/types.py

@@ -33,7 +33,9 @@ class CommandState(Enum):
     CANCELLED = "CANCELLED"
 
     @classmethod
-    def from_thrift_state(cls, state: ttypes.TOperationState) -> "CommandState":
+    def from_thrift_state(
+        cls, state: ttypes.TOperationState
+    ) -> Optional["CommandState"]:
         """
         Convert a Thrift TOperationState to a normalized CommandState.
 
@@ -75,7 +77,7 @@ def from_thrift_state(cls, state: ttypes.TOperationState) -> "CommandState":
         elif state == ttypes.TOperationState.CANCELED_STATE:
             return cls.CANCELLED
         else:
-            raise ValueError(f"Unknown command state: {state}")
+            return None
 
 
 class BackendType(Enum):

src/databricks/sql/result_set.py

@@ -1,15 +1,22 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional, Any, Union
+from typing import List, Optional, Any, Union, TYPE_CHECKING
 
 import logging
 import time
 import pandas
 
+from databricks.sql.backend.types import CommandId, CommandState
+
 try:
     import pyarrow
 except ImportError:
     pyarrow = None
 
+if TYPE_CHECKING:
+    from databricks.sql.backend.databricks_client import DatabricksClient
+    from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
+    from databricks.sql.client import Connection
+
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import Row
 from databricks.sql.exc import Error, RequestError, CursorAlreadyClosedError
@@ -25,10 +32,30 @@ class ResultSet(ABC):
     This class defines the interface that all concrete result set implementations must follow.
     """
 
-    def __init__(self, connection, backend, arraysize: int, buffer_size_bytes: int):
-        """Initialize the base ResultSet with common properties."""
+    def __init__(
+        self,
+        connection: "Connection",
+        backend: "DatabricksClient",
+        command_id: CommandId,
+        op_state: Optional[CommandState],
+        has_been_closed_server_side: bool,
+        arraysize: int,
+        buffer_size_bytes: int,
+    ):
+        """
+        A ResultSet manages the results of a single command.
+
+        :param connection: The parent connection that was used to execute this command
+        :param backend: The specialised backend client to be invoked in the fetch phase
+        :param execute_response: A `ExecuteResponse` class returned by a command execution
+        :param result_buffer_size_bytes: The size (in bytes) of the internal buffer + max fetch
+        amount :param arraysize: The max number of rows to fetch at a time (PEP-249)
+        """
+        self.command_id = command_id
+        self.op_state = op_state
+        self.has_been_closed_server_side = has_been_closed_server_side
         self.connection = connection
-        self.backend = backend  # Store the backend client directly
+        self.backend = backend
         self.arraysize = arraysize
         self.buffer_size_bytes = buffer_size_bytes
         self._next_row_index = 0
@@ -83,20 +110,36 @@ def fetchall_arrow(self) -> Any:
         """Fetch all remaining rows as an Arrow table."""
         pass
 
-    @abstractmethod
     def close(self) -> None:
-        """Close the result set and release any resources."""
-        pass
+        """
+        Close the result set.
+
+        If the connection has not been closed, and the result set has not already
+        been closed on the server for some other reason, issue a request to the server to close it.
+        """
+        try:
+            if (
+                self.op_state != CommandState.CLOSED
+                and not self.has_been_closed_server_side
+                and self.connection.open
+            ):
+                self.backend.close_command(self.command_id)
+        except RequestError as e:
+            if isinstance(e.args[1], CursorAlreadyClosedError):
+                logger.info("Operation was canceled by a prior request")
+        finally:
+            self.has_been_closed_server_side = True
+            self.op_state = CommandState.CLOSED
 
 
 class ThriftResultSet(ResultSet):
     """ResultSet implementation for the Thrift backend."""
 
     def __init__(
         self,
-        connection,
+        connection: "Connection",
         execute_response: ExecuteResponse,
-        thrift_client,  # Pass the specific ThriftDatabricksClient instance
+        thrift_client: "ThriftDatabricksClient",
         buffer_size_bytes: int = 104857600,
         arraysize: int = 10000,
         use_cloud_fetch: bool = True,
@@ -112,11 +155,20 @@ def __init__(
             arraysize: Default number of rows to fetch
             use_cloud_fetch: Whether to use cloud fetch for retrieving results
         """
-        super().__init__(connection, thrift_client, arraysize, buffer_size_bytes)
+        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,
+            arraysize,
+            buffer_size_bytes,
+        )
 
         # Initialize ThriftResultSet-specific attributes
-        self.command_id = execute_response.command_id
-        self.op_state = execute_response.status
         self.has_been_closed_server_side = execute_response.has_been_closed_server_side
         self.has_more_rows = execute_response.has_more_rows
         self.lz4_compressed = execute_response.lz4_compressed
@@ -127,11 +179,15 @@ def __init__(
 
         # Initialize results queue
         if execute_response.arrow_queue:
+            # In this case the server has taken the fast path and returned an initial batch of
+            # results
             self.results = execute_response.arrow_queue
         else:
+            # In this case, there are results waiting on the server so we fetch now for simplicity
             self._fill_results_buffer()
 
     def _fill_results_buffer(self):
+        # At initialization or if the server does not have cloud fetch result links available
         results, has_more_rows = self.backend.fetch_results(
             command_id=self.command_id,
             max_rows=self.arraysize,
@@ -336,28 +392,24 @@ def fetchmany(self, size: int) -> List[Row]:
         else:
             return self._convert_arrow_table(self.fetchmany_arrow(size))
 
-    def close(self) -> None:
-        """
-        Close the cursor.
-
-        If the connection has not been closed, and the cursor has not already
-        been closed on the server for some other reason, issue a request to the server to close it.
-        """
-        try:
-            if (
-                self.op_state != ttypes.TOperationState.CLOSED_STATE
-                and not self.has_been_closed_server_side
-                and self.connection.open
-            ):
-                self.backend.close_command(self.command_id)
-        except RequestError as e:
-            if isinstance(e.args[1], CursorAlreadyClosedError):
-                logger.info("Operation was canceled by a prior request")
-        finally:
-            self.has_been_closed_server_side = True
-            self.op_state = ttypes.TOperationState.CLOSED_STATE
-
     @property
     def is_staging_operation(self) -> bool:
         """Whether this result set represents a staging operation."""
         return self._is_staging_operation
+
+    @staticmethod
+    def _get_schema_description(table_schema_message):
+        """
+        Takes a TableSchema message and returns a description 7-tuple as specified by PEP-249
+        """
+
+        def map_col_type(type_):
+            if type_.startswith("decimal"):
+                return "decimal"
+            else:
+                return type_
+
+        return [
+            (column.name, map_col_type(column.datatype), None, None, None, None, None)
+            for column in table_schema_message.columns
+        ]