Merge branch 'sea-migration' into sea-http-client

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)
 
@@ -188,7 +188,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/queue.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union, TYPE_CHECKING
 
 from databricks.sql.cloudfetch.download_manager import ResultFileDownloadManager
 
@@ -12,12 +12,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
@@ -109,7 +110,7 @@ def __init__(
         result_data: ResultData,
         max_download_threads: int,
         ssl_options: SSLOptions,
-        sea_client: "SeaDatabricksClient",
+        sea_client: SeaDatabricksClient,
         statement_id: str,
         total_chunk_count: int,
         lz4_compressed: bool = False,
@@ -140,6 +141,7 @@ def __init__(
 
         self._sea_client = sea_client
         self._statement_id = statement_id
+        self._total_chunk_count = total_chunk_count
 
         logger.debug(
             "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
@@ -154,12 +156,11 @@ def __init__(
             return None
 
         # Track the current chunk we're processing
-        self._current_chunk_link = first_link
-
+        self._current_chunk_index = 0
         # Initialize table and position
-        self.table = self._create_table_from_link(self._current_chunk_link)
+        self.table = self._create_table_from_link(first_link)
 
-    def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink:
+    def _convert_to_thrift_link(self, link: ExternalLink) -> TSparkArrowResultLink:
         """Convert SEA external links to Thrift format for compatibility with existing download manager."""
         # Parse the ISO format expiration time
         expiry_time = int(dateutil.parser.parse(link.expiration).timestamp())
@@ -172,36 +173,24 @@ def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink
             httpHeaders=link.http_headers or {},
         )
 
-    def _progress_chunk_link(self):
+    def _get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]:
         """Progress to the next chunk link."""
-        if not self._current_chunk_link:
-            return None
-
-        next_chunk_index = self._current_chunk_link.next_chunk_index
-
-        if next_chunk_index is None:
-            self._current_chunk_link = None
+        if chunk_index >= self._total_chunk_count:
             return None
 
         try:
-            self._current_chunk_link = self._sea_client.get_chunk_link(
-                self._statement_id, next_chunk_index
-            )
+            return self._sea_client.get_chunk_link(self._statement_id, chunk_index)
         except Exception as e:
             raise ServerOperationError(
-                f"Error fetching link for chunk {next_chunk_index}: {e}",
+                f"Error fetching link for chunk {chunk_index}: {e}",
                 {
                     "operation-id": self._statement_id,
                     "diagnostic-info": None,
                 },
             )
 
-        logger.debug(
-            f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_link}"
-        )
-
     def _create_table_from_link(
-        self, link: "ExternalLink"
+        self, link: ExternalLink
     ) -> Union["pyarrow.Table", None]:
         """Create a table from a link."""
 
@@ -215,11 +204,8 @@ def _create_table_from_link(
 
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
-
-        self._progress_chunk_link()
-
-        if not self._current_chunk_link:
-            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
+        self._current_chunk_index += 1
+        next_chunk_link = self._get_chunk_link(self._current_chunk_index)
+        if not next_chunk_link:
             return None
-
-        return self._create_table_from_link(self._current_chunk_link)
+        return self._create_table_from_link(next_chunk_link)

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