Jun 12, 2025 · Jun 9, 2025 · Jun 9, 2025 · Jun 9, 2025 · Jun 9, 2025 · Jun 9, 2025
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
        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.
        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:
            :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
            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
            (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
 """
 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()
Original file line number	Diff line number	Diff line change
Expand Up		@@ -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.
Expand DownExpand Up		@@ -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:
Expand DownExpand Up		@@ -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
Expand DownExpand Up		@@ -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
Expand DownExpand Up		@@ -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")
Original file line number	Diff line number	Diff line change
		@@ -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()