Merge branch 'ext-links-sea' into sea-hybrid

Author: varun-edachali-dbx

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
 
@@ -14,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
@@ -128,7 +129,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,
@@ -160,6 +161,7 @@ def __init__(
         self._sea_client = sea_client
         self._statement_id = statement_id
         self._total_chunk_count = total_chunk_count
+        self._total_chunk_count = total_chunk_count
 
         logger.debug(
             "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
@@ -175,19 +177,14 @@ def __init__(
         first_link = self._chunk_index_to_link.get(self._current_chunk_index, None)
         if not first_link:
             # possibly an empty response
-            return
-
-        self.download_manager = ResultFileDownloadManager(
-            links=[],
-            max_download_threads=max_download_threads,
-            lz4_compressed=lz4_compressed,
-            ssl_options=ssl_options,
-        )
+            return None
 
+        # Track the current chunk we're processing
+        self._current_chunk_index = 0
         # Initialize table and position
         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())
@@ -220,7 +217,7 @@ def _get_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
         return link
 
     def _create_table_from_link(
-        self, link: "ExternalLink"
+        self, link: ExternalLink
     ) -> Union["pyarrow.Table", None]:
         """Create a table from a link."""
 

tests/unit/test_sea_queue.py

@@ -364,7 +364,7 @@ def test_init_with_valid_initial_link(
 
         # Verify attributes
         assert queue._statement_id == "test-statement-123"
-        assert queue._current_chunk_link == sample_external_link
+        assert queue._current_chunk_index == 0
 
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
     @patch("databricks.sql.backend.sea.queue.logger")
@@ -390,91 +390,31 @@ def test_init_no_initial_links(
         )
         assert queue.table is None
 
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_progress_chunk_link_no_current_link(self, mock_logger):
-        """Test _progress_chunk_link with no current link."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = None
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._progress_chunk_link(queue)
-
-        # Verify the result is None
-        assert result is None
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_progress_chunk_link_no_next_chunk(self, mock_logger):
-        """Test _progress_chunk_link with no next chunk index."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = ExternalLink(
-            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=None,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._progress_chunk_link(queue)
-
-        # Verify the result is None
-        assert result is None
-        assert queue._current_chunk_link is None
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_create_next_table_no_current_link(self, mock_logger):
-        """Test _create_next_table with no current link."""
-        # Create a queue instance without initializing
-        queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = None
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._create_next_table(queue)
-
-        # Verify debug message was logged
-        mock_logger.debug.assert_called_with(
-            "SeaCloudFetchQueue: No current chunk link, returning"
-        )
-
-        # Verify the result is None
-        assert result is None
-
     @patch("databricks.sql.backend.sea.queue.logger")
     def test_create_next_table_success(self, mock_logger):
         """Test _create_next_table with successful table creation."""
         # Create a queue instance without initializing
         queue = Mock(spec=SeaCloudFetchQueue)
-        queue._current_chunk_link = ExternalLink(
-            external_link="https://example.com/data/chunk0",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=50,
-            chunk_index=0,
-            next_chunk_index=1,
-            http_headers={"Authorization": "Bearer token123"},
-        )
+        queue._current_chunk_index = 0
         queue.download_manager = Mock()
 
         # Mock the dependencies
         mock_table = Mock()
-        queue._create_table_at_offset = Mock(return_value=mock_table)
+        mock_chunk_link = Mock()
+        queue._get_chunk_link = Mock(return_value=mock_chunk_link)
         queue._create_table_from_link = Mock(return_value=mock_table)
-        queue._progress_chunk_link = Mock()
 
         # Call the method directly
         result = SeaCloudFetchQueue._create_next_table(queue)
 
-        # Verify the table was created
-        queue._create_table_from_link.assert_called_once_with(queue._current_chunk_link)
+        # Verify the chunk index was incremented
+        assert queue._current_chunk_index == 1
+
+        # Verify the chunk link was retrieved
+        queue._get_chunk_link.assert_called_once_with(1)
 
-        # Verify progress was called
-        queue._progress_chunk_link.assert_called_once()
+        # Verify the table was created from the link
+        queue._create_table_from_link.assert_called_once_with(mock_chunk_link)
 
         # Verify the result is the table
         assert result == mock_table