diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index cf6940bb2..38b8a3c2f 100644
--- a/src/databricks/sql/result_set.py
+++ b/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")
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
new file mode 100644
index 000000000..c596dbc14
--- /dev/null
+++ b/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()
pFad - Phonifier reborn
Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.
Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.
Alternative Proxies:
Alternative Proxy
pFad Proxy
pFad v3 Proxy
pFad v4 Proxy