simplify download process: no pre-fetching

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

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

@@ -170,7 +170,6 @@ def __init__(
 
         # Track the current chunk we're processing
         self._current_chunk_link: Optional["ExternalLink"] = initial_link
-        self._download_current_link()
 
         # Initialize table and position
         self.table = self._create_next_table()
@@ -188,18 +187,6 @@ def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink
             httpHeaders=link.http_headers or {},
         )
 
-    def _download_current_link(self):
-        """Download the current chunk link."""
-        if not self._current_chunk_link:
-            return None
-
-        if not self.download_manager:
-            logger.debug("SeaCloudFetchQueue: No download manager, returning")
-            return None
-
-        thrift_link = self._convert_to_thrift_link(self._current_chunk_link)
-        self.download_manager.add_link(thrift_link)
-
     def _progress_chunk_link(self):
         """Progress to the next chunk link."""
         if not self._current_chunk_link:
@@ -221,19 +208,26 @@ def _progress_chunk_link(self):
                     next_chunk_index, e
                 )
             )
+            self._current_chunk_link = None
             return None
 
         logger.debug(
             f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_link}"
         )
-        self._download_current_link()
 
     def _create_next_table(self) -> Union["pyarrow.Table", None]:
         """Create next table by retrieving the logical next downloaded file."""
         if not self._current_chunk_link:
             logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
             return None
 
+        if not self.download_manager:
+            logger.debug("SeaCloudFetchQueue: No download manager, returning")
+            return None
+
+        thrift_link = self._convert_to_thrift_link(self._current_chunk_link)
+        self.download_manager.add_link(thrift_link)
+
         row_offset = self._current_chunk_link.row_offset
         arrow_table = self._create_table_at_offset(row_offset)
 

src/databricks/sql/utils.py

@@ -237,7 +237,7 @@ def __init__(
         self.table = None
         self.table_row_index = 0
 
-        # Initialize download manager - will be set by subclasses
+        # Initialize download manager
         self.download_manager: Optional["ResultFileDownloadManager"] = None
 
     def remaining_rows(self) -> "pyarrow.Table":

tests/unit/test_sea_queue.py

@@ -487,76 +487,6 @@ def test_init_non_zero_chunk_index(
         # Verify download manager wasn't created (no chunk 0)
         mock_download_manager_class.assert_not_called()
 
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_download_current_link_no_current_link(self, mock_logger):
-        """Test _download_current_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._download_current_link(queue)
-
-        # Verify the result is None
-        assert result is None
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_download_current_link_no_download_manager(
-        self, mock_logger, mock_sea_client, ssl_options
-    ):
-        """Test _download_current_link with no download manager."""
-        # 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=1,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-        queue.download_manager = None
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._download_current_link(queue)
-
-        # Verify debug message was logged
-        mock_logger.debug.assert_called_with(
-            "SeaCloudFetchQueue: No download manager, returning"
-        )
-
-        # Verify the result is None
-        assert result is None
-
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_download_current_link_success(self, mock_logger):
-        """Test _download_current_link with successful download."""
-        # 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=1,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-        queue.download_manager = Mock()
-
-        # Mock the _convert_to_thrift_link method
-        mock_thrift_link = Mock()
-        queue._convert_to_thrift_link = Mock(return_value=mock_thrift_link)
-
-        # Call the method directly
-        SeaCloudFetchQueue._download_current_link(queue)
-
-        # Verify the download manager add_link was called
-        queue.download_manager.add_link.assert_called_once_with(mock_thrift_link)
-
     @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."""
@@ -610,7 +540,6 @@ def test_progress_chunk_link_success(self, mock_logger, mock_sea_client):
         )
         queue._sea_client = mock_sea_client
         queue._statement_id = "test-statement-123"
-        queue._download_current_link = Mock()
 
         # Setup the mock client to return a new link
         next_link = ExternalLink(
@@ -636,9 +565,6 @@ def test_progress_chunk_link_success(self, mock_logger, mock_sea_client):
             f"SeaCloudFetchQueue: Progressed to link for chunk 1: {next_link}"
         )
 
-        # Verify _download_current_link was called
-        queue._download_current_link.assert_called_once()
-
     @patch("databricks.sql.backend.sea.queue.logger")
     def test_progress_chunk_link_error(self, mock_logger, mock_sea_client):
         """Test _progress_chunk_link with error during chunk fetch."""
@@ -710,6 +636,7 @@ def test_create_next_table_success(self, mock_logger):
             next_chunk_index=1,
             http_headers={"Authorization": "Bearer token123"},
         )
+        queue.download_manager = Mock()
 
         # Mock the dependencies
         mock_table = Mock()

tests/unit/test_thrift_field_ids.py

@@ -16,72 +16,77 @@ class TestThriftFieldIds:
 
     # Known exceptions that exceed the field ID limit
     KNOWN_EXCEPTIONS = {
-        ('TExecuteStatementReq', 'enforceEmbeddedSchemaCorrectness'): 3353,
-        ('TSessionHandle', 'serverProtocolVersion'): 3329,
+        ("TExecuteStatementReq", "enforceEmbeddedSchemaCorrectness"): 3353,
+        ("TSessionHandle", "serverProtocolVersion"): 3329,
     }
 
     def test_all_thrift_field_ids_are_within_allowed_range(self):
         """
         Validates that all field IDs in Thrift-generated classes are within the allowed range.
-        
+
         This test prevents field ID conflicts and ensures compatibility with different
         Thrift implementations and protocols.
         """
         violations = []
-        
+
         # Get all classes from the ttypes module
         for name, obj in inspect.getmembers(ttypes):
-            if (inspect.isclass(obj) and 
-                hasattr(obj, 'thrift_spec') and 
-                obj.thrift_spec is not None):
-                
+            if (
+                inspect.isclass(obj)
+                and hasattr(obj, "thrift_spec")
+                and obj.thrift_spec is not None
+            ):
+
                 self._check_class_field_ids(obj, name, violations)
-        
+
         if violations:
             error_message = self._build_error_message(violations)
             pytest.fail(error_message)
 
     def _check_class_field_ids(self, cls, class_name, violations):
         """
         Checks all field IDs in a Thrift class and reports violations.
-        
+
         Args:
             cls: The Thrift class to check
             class_name: Name of the class for error reporting
             violations: List to append violation messages to
         """
         thrift_spec = cls.thrift_spec
-        
+
         if not isinstance(thrift_spec, (tuple, list)):
             return
-            
+
         for spec_entry in thrift_spec:
             if spec_entry is None:
                 continue
-                
+
             # Thrift spec format: (field_id, field_type, field_name, ...)
             if isinstance(spec_entry, (tuple, list)) and len(spec_entry) >= 3:
                 field_id = spec_entry[0]
                 field_name = spec_entry[2]
-                
+
                 # Skip known exceptions
                 if (class_name, field_name) in self.KNOWN_EXCEPTIONS:
                     continue
-                
+
                 if isinstance(field_id, int) and field_id >= self.MAX_ALLOWED_FIELD_ID:
                     violations.append(
                         "{} field '{}' has field ID {} (exceeds maximum of {})".format(
-                            class_name, field_name, field_id, self.MAX_ALLOWED_FIELD_ID - 1
+                            class_name,
+                            field_name,
+                            field_id,
+                            self.MAX_ALLOWED_FIELD_ID - 1,
                         )
                     )
 
     def _build_error_message(self, violations):
         """
         Builds a comprehensive error message for field ID violations.
-        
+
         Args:
             violations: List of violation messages
-            
+
         Returns:
             Formatted error message
         """
@@ -90,8 +95,8 @@ def _build_error_message(self, violations):
             "This can cause compatibility issues and conflicts with reserved ID ranges.\n"
             "Violations found:\n".format(self.MAX_ALLOWED_FIELD_ID - 1)
         )
-        
+
         for violation in violations:
             error_message += "  - {}\n".format(violation)
-            
-        return error_message 
+
+        return error_message