fix typing, errors

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/backend.py

@@ -19,7 +19,7 @@
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
 
-from databricks.sql.backend.sea.result_set import SeaResultSet
+from databricks.sql.result_set import SeaResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (
@@ -29,7 +29,7 @@
     BackendType,
     ExecuteResponse,
 )
-from databricks.sql.exc import DatabaseError, ProgrammingError, ServerOperationError
+from databricks.sql.exc import DatabaseError, ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
 from databricks.sql.types import SSLOptions
 
@@ -135,7 +135,7 @@ def __init__(
         self.warehouse_id = self._extract_warehouse_id(http_path)
 
         # Initialize HTTP client
-        self.http_client = SeaHttpClient(
+        self._http_client = SeaHttpClient(
             server_hostname=server_hostname,
             port=port,
             http_path=http_path,
@@ -180,7 +180,7 @@ def _extract_warehouse_id(self, http_path: str) -> str:
             f"Note: SEA only works for warehouses."
         )
         logger.error(error_message)
-        raise ProgrammingError(error_message)
+        raise ValueError(error_message)
 
     @property
     def max_download_threads(self) -> int:
@@ -227,7 +227,7 @@ def open_session(
             schema=schema,
         )
 
-        response = self.http_client._make_request(
+        response = self._http_client._make_request(
             method="POST", path=self.SESSION_PATH, data=request_data.to_dict()
         )
 
@@ -252,22 +252,22 @@ def close_session(self, session_id: SessionId) -> None:
             session_id: The session identifier returned by open_session()
 
         Raises:
-            ProgrammingError: If the session ID is invalid
+            ValueError: If the session ID is invalid
             OperationalError: If there's an error closing the session
         """
 
         logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
 
         if session_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA session ID")
+            raise ValueError("Not a valid SEA session ID")
         sea_session_id = session_id.to_sea_session_id()
 
         request_data = DeleteSessionRequest(
             warehouse_id=self.warehouse_id,
             session_id=sea_session_id,
         )
 
-        self.http_client._make_request(
+        self._http_client._make_request(
             method="DELETE",
             path=self.SESSION_PATH_WITH_ID.format(sea_session_id),
             data=request_data.to_dict(),
@@ -462,7 +462,7 @@ def execute_command(
         """
 
         if session_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA session ID")
+            raise ValueError("Not a valid SEA session ID")
 
         sea_session_id = session_id.to_sea_session_id()
 
@@ -509,7 +509,7 @@ def execute_command(
             result_compression=result_compression,
         )
 
-        response_data = self.http_client._make_request(
+        response_data = self._http_client._make_request(
             method="POST", path=self.STATEMENT_PATH, data=request.to_dict()
         )
         response = ExecuteStatementResponse.from_dict(response_data)
@@ -546,16 +546,16 @@ def cancel_command(self, command_id: CommandId) -> None:
             command_id: Command identifier to cancel
 
         Raises:
-            ProgrammingError: If the command ID is invalid
+            ValueError: If the command ID is invalid
         """
 
         if command_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA command ID")
+            raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
 
         request = CancelStatementRequest(statement_id=sea_statement_id)
-        self.http_client._make_request(
+        self._http_client._make_request(
             method="POST",
             path=self.CANCEL_STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
@@ -569,16 +569,16 @@ def close_command(self, command_id: CommandId) -> None:
             command_id: Command identifier to close
 
         Raises:
-            ProgrammingError: If the command ID is invalid
+            ValueError: If the command ID is invalid
         """
 
         if command_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA command ID")
+            raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
 
         request = CloseStatementRequest(statement_id=sea_statement_id)
-        self.http_client._make_request(
+        self._http_client._make_request(
             method="DELETE",
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
@@ -595,7 +595,7 @@ def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
         sea_statement_id = command_id.to_sea_statement_id()
 
         request = GetStatementRequest(statement_id=sea_statement_id)
-        response_data = self.http_client._make_request(
+        response_data = self._http_client._make_request(
             method="GET",
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
@@ -615,7 +615,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             CommandState: The current state of the command
 
         Raises:
-            ProgrammingError: If the command ID is invalid
+            ValueError: If the command ID is invalid
         """
 
         response = self._poll_query(command_id)
@@ -643,27 +643,6 @@ def get_execution_result(
         response = self._poll_query(command_id)
         return self._response_to_result_set(response, cursor)
 
-    def get_chunk_links(
-        self, statement_id: str, chunk_index: int
-    ) -> List[ExternalLink]:
-        """
-        Get links for chunks starting from the specified index.
-        Args:
-            statement_id: The statement ID
-            chunk_index: The starting chunk index
-        Returns:
-            ExternalLink: External link for the chunk
-        """
-
-        response_data = self._http_client._make_request(
-            method="GET",
-            path=self.CHUNK_PATH_WITH_ID_AND_INDEX.format(statement_id, chunk_index),
-        )
-        response = GetChunksResponse.from_dict(response_data)
-
-        links = response.external_links or []
-        return links
-
     # == Metadata Operations ==
 
     def get_catalogs(

src/databricks/sql/result_set.py

@@ -1,11 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import List, Optional, Any, Union, Tuple, TYPE_CHECKING
+from typing import List, Optional, Any, TYPE_CHECKING
 
 import logging
-import time
 import pandas
 
-from databricks.sql.backend.sea.backend import SeaDatabricksClient
 
 try:
     import pyarrow
@@ -14,11 +12,12 @@
 
 if TYPE_CHECKING:
     from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
+    from databricks.sql.backend.sea.backend import SeaDatabricksClient
     from databricks.sql.client import Connection
+
 from databricks.sql.backend.databricks_client import DatabricksClient
-from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import Row
-from databricks.sql.exc import Error, RequestError, CursorAlreadyClosedError
+from databricks.sql.exc import RequestError, CursorAlreadyClosedError
 from databricks.sql.utils import ColumnTable, ColumnQueue
 from databricks.sql.backend.types import CommandId, CommandState, ExecuteResponse
 
@@ -136,7 +135,8 @@ def close(self) -> None:
         been closed on the server for some other reason, issue a request to the server to close it.
         """
         try:
-            self.results.close()
+            if self.results:
+                self.results.close()
             if (
                 self.status != CommandState.CLOSED
                 and not self.has_been_closed_server_side

tests/unit/test_sea_backend.py

@@ -132,7 +132,7 @@ def test_initialization(self, mock_http_client):
         assert client3.max_download_threads == 5
 
         # Test with invalid HTTP path
-        with pytest.raises(ProgrammingError) as excinfo:
+        with pytest.raises(ValueError) as excinfo:
             SeaDatabricksClient(
                 server_hostname="test-server.databricks.com",
                 port=443,
@@ -198,7 +198,7 @@ def test_session_management(self, sea_client, mock_http_client, thrift_session_i
         )
 
         # Test close_session with invalid ID type
-        with pytest.raises(ProgrammingError) as excinfo:
+        with pytest.raises(ValueError) as excinfo:
             sea_client.close_session(thrift_session_id)
         assert "Not a valid SEA session ID" in str(excinfo.value)
 
@@ -244,7 +244,7 @@ def test_command_execution_sync(
             assert result == "mock_result_set"
 
         # Test with invalid session ID
-        with pytest.raises(ProgrammingError) as excinfo:
+        with pytest.raises(ValueError) as excinfo:
             mock_thrift_handle = MagicMock()
             mock_thrift_handle.sessionId.guid = b"guid"
             mock_thrift_handle.sessionId.secret = b"secret"
@@ -449,7 +449,7 @@ def test_command_management(
         )
 
         # Test cancel_command with invalid ID
-        with pytest.raises(ProgrammingError) as excinfo:
+        with pytest.raises(ValueError) as excinfo:
             sea_client.cancel_command(thrift_command_id)
         assert "Not a valid SEA command ID" in str(excinfo.value)
 
@@ -463,7 +463,7 @@ def test_command_management(
         )
 
         # Test close_command with invalid ID
-        with pytest.raises(ProgrammingError) as excinfo:
+        with pytest.raises(ValueError) as excinfo:
             sea_client.close_command(thrift_command_id)
         assert "Not a valid SEA command ID" in str(excinfo.value)
 
@@ -522,7 +522,7 @@ def test_command_management(
         assert result.status == CommandState.SUCCEEDED
 
         # Test get_execution_result with invalid ID
-        with pytest.raises(ProgrammingError) as excinfo:
+        with pytest.raises(ValueError) as excinfo:
             sea_client.get_execution_result(thrift_command_id, mock_cursor)
         assert "Not a valid SEA command ID" in str(excinfo.value)
 
@@ -955,67 +955,3 @@ def test_get_columns(self, sea_client, sea_session_id, mock_cursor):
                     cursor=mock_cursor,
                 )
             assert "Catalog name is required for get_columns" in str(excinfo.value)
-
-    def test_get_chunk_links(self, sea_client, mock_http_client, sea_command_id):
-        """Test get_chunk_links method when links are available."""
-        # Setup mock response
-        mock_response = {
-            "external_links": [
-                {
-                    "external_link": "https://example.com/data/chunk0",
-                    "expiration": "2025-07-03T05:51:18.118009",
-                    "row_count": 100,
-                    "byte_count": 1024,
-                    "row_offset": 0,
-                    "chunk_index": 0,
-                    "next_chunk_index": 1,
-                    "http_headers": {"Authorization": "Bearer token123"},
-                }
-            ]
-        }
-        mock_http_client._make_request.return_value = mock_response
-
-        # Call the method
-        results = sea_client.get_chunk_links("test-statement-123", 0)
-
-        # Verify the HTTP client was called correctly
-        mock_http_client._make_request.assert_called_once_with(
-            method="GET",
-            path=sea_client.CHUNK_PATH_WITH_ID_AND_INDEX.format(
-                "test-statement-123", 0
-            ),
-        )
-
-        # Verify the results
-        assert isinstance(results, list)
-        assert len(results) == 1
-        result = results[0]
-        assert result.external_link == "https://example.com/data/chunk0"
-        assert result.expiration == "2025-07-03T05:51:18.118009"
-        assert result.row_count == 100
-        assert result.byte_count == 1024
-        assert result.row_offset == 0
-        assert result.chunk_index == 0
-        assert result.next_chunk_index == 1
-        assert result.http_headers == {"Authorization": "Bearer token123"}
-
-    def test_get_chunk_links_empty(self, sea_client, mock_http_client):
-        """Test get_chunk_links when no links are returned (empty list)."""
-        # Setup mock response with no matching chunk
-        mock_response = {"external_links": []}
-        mock_http_client._make_request.return_value = mock_response
-
-        # Call the method
-        results = sea_client.get_chunk_links("test-statement-123", 0)
-
-        # Verify the HTTP client was called correctly
-        mock_http_client._make_request.assert_called_once_with(
-            method="GET",
-            path=sea_client.CHUNK_PATH_WITH_ID_AND_INDEX.format(
-                "test-statement-123", 0
-            ),
-        )
-
-        # Verify the results are empty
-        assert isinstance(results, list)
-        assert results == []