Merge branch 'sea-migration' into sea-decouple-link-fetch

Author: varun-edachali-dbx

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

@@ -50,17 +50,17 @@
 
 
 def _filter_session_configuration(
-    session_configuration: Optional[Dict[str, str]]
-) -> Optional[Dict[str, str]]:
+    session_configuration: Optional[Dict[str, Any]],
+) -> Dict[str, str]:
     if not session_configuration:
-        return None
+        return {}
 
     filtered_session_configuration = {}
     ignored_configs: Set[str] = set()
 
     for key, value in session_configuration.items():
         if key.upper() in ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP:
-            filtered_session_configuration[key.lower()] = value
+            filtered_session_configuration[key.lower()] = str(value)
         else:
             ignored_configs.add(key)
 
@@ -190,7 +190,7 @@ def max_download_threads(self) -> int:
 
     def open_session(
         self,
-        session_configuration: Optional[Dict[str, str]],
+        session_configuration: Optional[Dict[str, Any]],
         catalog: Optional[str],
         schema: Optional[str],
     ) -> SessionId:

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

@@ -94,8 +94,7 @@ def _parse_result(data: Dict[str, Any]) -> ResultData:
 
     # Handle attachment field - decode from base64 if present
     attachment = result_data.get("attachment")
-    if attachment is not None and isinstance(attachment, str):
-        # Decode base64 string to bytes
+    if attachment is not None:
         attachment = base64.b64decode(attachment)
 
     return ResultData(

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

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC
-import threading
-from typing import Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, TYPE_CHECKING
 
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 
@@ -15,12 +14,13 @@
 
 import dateutil
 
-from databricks.sql.backend.sea.backend import SeaDatabricksClient
-from databricks.sql.backend.sea.models.base import (
-    ExternalLink,
-    ResultData,
-    ResultManifest,
-)
+if TYPE_CHECKING:
+    from databricks.sql.backend.sea.backend import SeaDatabricksClient
+    from databricks.sql.backend.sea.models.base import (
+        ExternalLink,
+        ResultData,
+        ResultManifest,
+    )
 from databricks.sql.backend.sea.utils.constants import ResultFormat
 from databricks.sql.exc import ProgrammingError, ServerOperationError
 from databricks.sql.thrift_api.TCLIService.ttypes import TSparkArrowResultLink

tests/unit/test_sea_backend.py

@@ -624,6 +624,71 @@ def test_utility_methods(self, sea_client):
         assert description[1][1] == "INT"  # type_code
         assert description[1][6] is False  # null_ok
 
+    def test_filter_session_configuration(self):
+        """Test that _filter_session_configuration converts all values to strings."""
+        session_config = {
+            "ANSI_MODE": True,
+            "statement_timeout": 3600,
+            "TIMEZONE": "UTC",
+            "enable_photon": False,
+            "MAX_FILE_PARTITION_BYTES": 128.5,
+            "unsupported_param": "value",
+            "ANOTHER_UNSUPPORTED": 42,
+        }
+
+        result = _filter_session_configuration(session_config)
+
+        # Verify result is not None
+        assert result is not None
+
+        # Verify all returned values are strings
+        for key, value in result.items():
+            assert isinstance(
+                value, str
+            ), f"Value for key '{key}' is not a string: {type(value)}"
+
+        # Verify specific conversions
+        expected_result = {
+            "ansi_mode": "True",  # boolean True -> "True", key lowercased
+            "statement_timeout": "3600",  # int -> "3600", key lowercased
+            "timezone": "UTC",  # string -> "UTC", key lowercased
+            "enable_photon": "False",  # boolean False -> "False", key lowercased
+            "max_file_partition_bytes": "128.5",  # float -> "128.5", key lowercased
+        }
+
+        assert result == expected_result
+
+        # Test with None input
+        assert _filter_session_configuration(None) == {}
+
+        # Test with only unsupported parameters
+        unsupported_config = {
+            "unsupported_param1": "value1",
+            "unsupported_param2": 123,
+        }
+        result = _filter_session_configuration(unsupported_config)
+        assert result == {}
+
+        # Test case insensitivity for keys
+        case_insensitive_config = {
+            "ansi_mode": "false",  # lowercase key
+            "STATEMENT_TIMEOUT": 7200,  # uppercase key
+            "TiMeZoNe": "America/New_York",  # mixed case key
+        }
+        result = _filter_session_configuration(case_insensitive_config)
+        expected_case_result = {
+            "ansi_mode": "false",
+            "statement_timeout": "7200",
+            "timezone": "America/New_York",
+        }
+        assert result == expected_case_result
+
+        # Verify all values are strings in case insensitive test
+        for key, value in result.items():
+            assert isinstance(
+                value, str
+            ), f"Value for key '{key}' is not a string: {type(value)}"
+
     def test_results_message_to_execute_response_is_staging_operation(self, sea_client):
         """Test that is_staging_operation is correctly set from manifest.is_volume_operation."""
         # Test when is_volume_operation is True
@@ -893,3 +958,67 @@ 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 == []