Introduce preliminary SEA Result Set

src/databricks/sql/result_set.py

@@ -45,6 +45,8 @@ def __init__(
         results_queue=None,
         description=None,
         is_staging_operation: bool = False,
+        lz4_compressed: bool = False,
+        arrow_schema_bytes: Optional[bytes] = None,
     ):
         """
         A ResultSet manages the results of a single command.
@@ -75,6 +77,8 @@ def __init__(
         self.is_direct_results = is_direct_results
         self.results = results_queue
         self._is_staging_operation = is_staging_operation
+        self.lz4_compressed = lz4_compressed
+        self._arrow_schema_bytes = arrow_schema_bytes
 
     def __iter__(self):
         while True:
@@ -177,10 +181,10 @@ def __init__(
             :param ssl_options: SSL options for cloud fetch
             :param is_direct_results: Whether there are more rows to fetch
         """
+
         # Initialize ThriftResultSet-specific attributes
-        self._arrow_schema_bytes = execute_response.arrow_schema_bytes
         self._use_cloud_fetch = use_cloud_fetch
-        self.lz4_compressed = execute_response.lz4_compressed
+        self.is_direct_results = is_direct_results
 
         # Build the results queue if t_row_set is provided
         results_queue = None
@@ -211,6 +215,8 @@ def __init__(
             results_queue=results_queue,
             description=execute_response.description,
             is_staging_operation=execute_response.is_staging_operation,
+            lz4_compressed=execute_response.lz4_compressed,
+            arrow_schema_bytes=execute_response.arrow_schema_bytes,
         )
 
         # Initialize results queue if not provided
@@ -438,3 +444,82 @@ def map_col_type(type_):
             (column.name, map_col_type(column.datatype), None, None, None, None, None)
             for column in table_schema_message.columns
         ]
+
+
+class SeaResultSet(ResultSet):
+    """ResultSet implementation for SEA backend."""
+
+    def __init__(
+        self,
+        connection: "Connection",
+        execute_response: "ExecuteResponse",
+        sea_client: "SeaDatabricksClient",
+        buffer_size_bytes: int = 104857600,
+        arraysize: int = 10000,
+        result_data=None,
+        manifest=None,
+    ):
+        """
+        Initialize a SeaResultSet with the response from a SEA query execution.
+
+        Args:
+            connection: The parent connection
+            execute_response: Response from the execute command
+            sea_client: The SeaDatabricksClient instance for direct access
+            buffer_size_bytes: Buffer size for fetching results
+            arraysize: Default number of rows to fetch
+            result_data: Result data from SEA response (optional)
+            manifest: Manifest from SEA response (optional)
+        """
+
+        super().__init__(
+            connection=connection,
+            backend=sea_client,
+            arraysize=arraysize,
+            buffer_size_bytes=buffer_size_bytes,
+            command_id=execute_response.command_id,
+            status=execute_response.status,
+            has_been_closed_server_side=execute_response.has_been_closed_server_side,
+            description=execute_response.description,
+            is_staging_operation=execute_response.is_staging_operation,
+            lz4_compressed=execute_response.lz4_compressed,
+            arrow_schema_bytes=execute_response.arrow_schema_bytes,
+        )
+
+    def _fill_results_buffer(self):
+        """Fill the results buffer from the backend."""
+        raise NotImplementedError(
+            "_fill_results_buffer is not implemented for SEA backend"
+        )
+
+    def fetchone(self) -> Optional[Row]:
+        """
+        Fetch the next row of a query result set, returning a single sequence,
+        or None when no more data is available.
+        """
+
+        raise NotImplementedError("fetchone is not implemented for SEA backend")
+
+    def fetchmany(self, size: Optional[int] = None) -> List[Row]:
+        """
+        Fetch the next set of rows of a query result, returning a list of rows.
+
+        An empty sequence is returned when no more rows are available.
+        """
+
+        raise NotImplementedError("fetchmany is not implemented for SEA backend")
+
+    def fetchall(self) -> List[Row]:
+        """
+        Fetch all (remaining) rows of a query result, returning them as a list of rows.
+        """
+
+        raise NotImplementedError("fetchall is not implemented for SEA backend")
+
+    def fetchmany_arrow(self, size: int) -> Any:
+        """Fetch the next set of rows as an Arrow table."""
+        raise NotImplementedError("fetchmany_arrow is not implemented for SEA backend")
+
+    def fetchall_arrow(self) -> Any:
+        """Fetch all remaining rows as an Arrow table."""
+        raise NotImplementedError("fetchall_arrow is not implemented for SEA backend")

tests/unit/test_sea_result_set.py

@@ -0,0 +1,201 @@
+"""
+Tests for the SeaResultSet class.
+
+This module contains tests for the SeaResultSet class, which implements
+the result set functionality for the SEA (Statement Execution API) backend.
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock, Mock
+
+from databricks.sql.result_set import SeaResultSet
+from databricks.sql.backend.types import CommandId, CommandState, BackendType
+
+
+class TestSeaResultSet:
+    """Test suite for the SeaResultSet class."""
+
+    @pytest.fixture
+    def mock_connection(self):
+        """Create a mock connection."""
+        connection = Mock()
+        connection.open = True
+        return connection
+
+    @pytest.fixture
+    def mock_sea_client(self):
+        """Create a mock SEA client."""
+        return Mock()
+
+    @pytest.fixture
+    def execute_response(self):
+        """Create a sample execute response."""
+        mock_response = Mock()
+        mock_response.command_id = CommandId.from_sea_statement_id("test-statement-123")
+        mock_response.status = CommandState.SUCCEEDED
+        mock_response.has_been_closed_server_side = False
+        mock_response.is_direct_results = False
+        mock_response.results_queue = None
+        mock_response.description = [
+            ("test_value", "INT", None, None, None, None, None)
+        ]
+        mock_response.is_staging_operation = False
+        return mock_response
+
+    def test_init_with_execute_response(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test initializing SeaResultSet with an execute response."""
+        result_set = SeaResultSet(
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
+        )
+
+        # Verify basic properties
+        assert result_set.command_id == execute_response.command_id
+        assert result_set.status == CommandState.SUCCEEDED
+        assert result_set.connection == mock_connection
+        assert result_set.backend == mock_sea_client
+        assert result_set.buffer_size_bytes == 1000
+        assert result_set.arraysize == 100
+        assert result_set.description == execute_response.description
+
+    def test_close(self, mock_connection, mock_sea_client, execute_response):
+        """Test closing a result set."""
+        result_set = SeaResultSet(
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
+        )
+
+        # Close the result set
+        result_set.close()
+
+        # Verify the backend's close_command was called
+        mock_sea_client.close_command.assert_called_once_with(result_set.command_id)
+        assert result_set.has_been_closed_server_side is True
+        assert result_set.status == CommandState.CLOSED
+
+    def test_close_when_already_closed_server_side(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test closing a result set that has already been closed server-side."""
+        result_set = SeaResultSet(
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
+        )
+        result_set.has_been_closed_server_side = True
+
+        # Close the result set
+        result_set.close()
+
+        # Verify the backend's close_command was NOT called
+        mock_sea_client.close_command.assert_not_called()
+        assert result_set.has_been_closed_server_side is True
+        assert result_set.status == CommandState.CLOSED
+
+    def test_close_when_connection_closed(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test closing a result set when the connection is closed."""
+        mock_connection.open = False
+        result_set = SeaResultSet(
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
+        )
+
+        # Close the result set
+        result_set.close()
+
+        # Verify the backend's close_command was NOT called
+        mock_sea_client.close_command.assert_not_called()
+        assert result_set.has_been_closed_server_side is True
+        assert result_set.status == CommandState.CLOSED
+
+    def test_unimplemented_methods(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test that unimplemented methods raise NotImplementedError."""
+        result_set = SeaResultSet(
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
+        )
+
+        # Test each unimplemented method individually with specific error messages
+        with pytest.raises(
+            NotImplementedError, match="fetchone is not implemented for SEA backend"
+        ):
+            result_set.fetchone()
+
+        with pytest.raises(
+            NotImplementedError, match="fetchmany is not implemented for SEA backend"
+        ):
+            result_set.fetchmany(10)
+
+        with pytest.raises(
+            NotImplementedError, match="fetchmany is not implemented for SEA backend"
+        ):
+            # Test with default parameter value
+            result_set.fetchmany()
+
+        with pytest.raises(
+            NotImplementedError, match="fetchall is not implemented for SEA backend"
+        ):
+            result_set.fetchall()
+
+        with pytest.raises(
+            NotImplementedError,
+            match="fetchmany_arrow is not implemented for SEA backend",
+        ):
+            result_set.fetchmany_arrow(10)
+
+        with pytest.raises(
+            NotImplementedError,
+            match="fetchall_arrow is not implemented for SEA backend",
+        ):
+            result_set.fetchall_arrow()
+
+        with pytest.raises(
+            NotImplementedError, match="fetchone is not implemented for SEA backend"
+        ):
+            # Test iteration protocol (calls fetchone internally)
+            next(iter(result_set))
+
+        with pytest.raises(
+            NotImplementedError, match="fetchone is not implemented for SEA backend"
+        ):
+            # Test using the result set in a for loop
+            for row in result_set:
+                pass
+
+    def test_fill_results_buffer_not_implemented(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Test that _fill_results_buffer raises NotImplementedError."""
+        result_set = SeaResultSet(
+            connection=mock_connection,
+            execute_response=execute_response,
+            sea_client=mock_sea_client,
+            buffer_size_bytes=1000,
+            arraysize=100,
+        )
+
+        with pytest.raises(
+            NotImplementedError,
+            match="_fill_results_buffer is not implemented for SEA backend",
+        ):
+            result_set._fill_results_buffer()