resolve merge artifacts

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

Author: varun-edachali-dbx

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

@@ -281,51 +281,14 @@ def __init__(
         self.link_fetcher.start()
 
         # Initialize table and position
-        self.table = self._create_table_from_link(self._current_chunk_link)
-
-    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())
-        return TSparkArrowResultLink(
-            fileLink=link.external_link,
-            expiryTime=expiry_time,
-            rowCount=link.row_count,
-            bytesNum=link.byte_count,
-            startRowOffset=link.row_offset,
-            httpHeaders=link.http_headers or {},
-        )
+        self.table = self._create_next_table()
 
     def _get_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
         if chunk_index not in self._chunk_index_to_link:
             links = self._sea_client.get_chunk_links(self._statement_id, chunk_index)
             self._chunk_index_to_link.update({link.chunk_index: link for link in links})
         return self._chunk_index_to_link.get(chunk_index, None)
 
-    def _progress_chunk_link(self):
-        """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
-            return None
-
-        self._current_chunk_link = self._get_chunk_link(next_chunk_index)
-        if not self._current_chunk_link:
-            logger.error(
-                "SeaCloudFetchQueue: unable to retrieve link for chunk {}".format(
-                    next_chunk_index
-                )
-            )
-            return None
-
-        logger.debug(
-            f"SeaCloudFetchQueue: Progressed to link for chunk {next_chunk_index}: {self._current_chunk_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:
@@ -345,8 +308,4 @@ def _create_next_table(self) -> Union["pyarrow.Table", None]:
 
         self.current_chunk_index += 1
 
-        if not self._current_chunk_link:
-            logger.debug("SeaCloudFetchQueue: No current chunk link, returning")
-            return None
-
-        return self._create_table_from_link(self._current_chunk_link)
+        return arrow_table

tests/unit/test_sea_queue.py

@@ -213,9 +213,7 @@ def test_build_queue_arrow_stream(
 
         with patch(
             "databricks.sql.backend.sea.queue.ResultFileDownloadManager"
-        ), patch.object(
-            SeaCloudFetchQueue, "_create_table_from_link", return_value=None
-        ):
+        ), patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
             queue = SeaResultSetQueueFactory.build_queue(
                 result_data=result_data,
                 manifest=arrow_manifest,
@@ -298,79 +296,6 @@ def sample_external_link_no_headers(self):
             http_headers=None,
         )
 
-    def test_convert_to_thrift_link(self, sample_external_link):
-        """Test conversion of ExternalLink to TSparkArrowResultLink."""
-        queue = Mock(spec=SeaCloudFetchQueue)
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._convert_to_thrift_link(queue, sample_external_link)
-
-        # Verify the conversion
-        assert result.fileLink == sample_external_link.external_link
-        assert result.rowCount == sample_external_link.row_count
-        assert result.bytesNum == sample_external_link.byte_count
-        assert result.startRowOffset == sample_external_link.row_offset
-        assert result.httpHeaders == sample_external_link.http_headers
-
-    def test_convert_to_thrift_link_no_headers(self, sample_external_link_no_headers):
-        """Test conversion of ExternalLink with no headers to TSparkArrowResultLink."""
-        queue = Mock(spec=SeaCloudFetchQueue)
-
-        # Call the method directly
-        result = SeaCloudFetchQueue._convert_to_thrift_link(
-            queue, sample_external_link_no_headers
-        )
-
-        # Verify the conversion
-        assert result.fileLink == sample_external_link_no_headers.external_link
-        assert result.rowCount == sample_external_link_no_headers.row_count
-        assert result.bytesNum == sample_external_link_no_headers.byte_count
-        assert result.startRowOffset == sample_external_link_no_headers.row_offset
-        assert result.httpHeaders == {}
-
-    @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_init_with_valid_initial_link(
-        self,
-        mock_logger,
-        mock_download_manager_class,
-        mock_sea_client,
-        ssl_options,
-        description,
-        sample_external_link,
-    ):
-        """Test initialization with valid initial link."""
-        mock_download_manager = Mock()
-        mock_download_manager_class.return_value = mock_download_manager
-
-        # Create a queue with valid initial link
-        with patch.object(SeaCloudFetchQueue, "_create_next_table", return_value=None):
-            queue = SeaCloudFetchQueue(
-                initial_links=[sample_external_link],
-                max_download_threads=5,
-                ssl_options=ssl_options,
-                sea_client=mock_sea_client,
-                statement_id="test-statement-123",
-                total_chunk_count=1,
-                lz4_compressed=False,
-                description=description,
-            )
-
-        # Verify debug message was logged
-        mock_logger.debug.assert_called_with(
-            "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
-                "test-statement-123", 1
-            )
-        )
-
-        # Verify download manager was created
-        mock_download_manager_class.assert_called_once()
-
-        # Verify attributes
-        assert queue._statement_id == "test-statement-123"
-        assert queue._current_chunk_link == sample_external_link
-        assert queue.download_manager == mock_download_manager
-
     @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
     @patch("databricks.sql.backend.sea.queue.logger")
     def test_init_no_initial_links(
@@ -410,136 +335,3 @@ def test_init_no_initial_links(
             not hasattr(queue, "_current_chunk_link")
             or queue._current_chunk_link is None
         )
-
-    @patch("databricks.sql.backend.sea.queue.ResultFileDownloadManager")
-    @patch("databricks.sql.backend.sea.queue.logger")
-    def test_init_non_zero_chunk_index(
-        self,
-        mock_logger,
-        mock_download_manager_class,
-        mock_sea_client,
-        ssl_options,
-        description,
-    ):
-        """Test initialization with non-zero chunk index initial link."""
-        # Create a link with chunk_index != 0
-        non_zero_link = ExternalLink(
-            external_link="https://example.com/data/chunk1",
-            expiration="2025-07-03T05:51:18.118009",
-            row_count=100,
-            byte_count=1024,
-            row_offset=100,
-            chunk_index=1,
-            next_chunk_index=2,
-            http_headers={"Authorization": "Bearer token123"},
-        )
-
-        # Create a queue with non-zero chunk index
-        queue = SeaCloudFetchQueue(
-            initial_links=[non_zero_link],
-            max_download_threads=5,
-            ssl_options=ssl_options,
-            sea_client=mock_sea_client,
-            statement_id="test-statement-123",
-            total_chunk_count=1,
-            lz4_compressed=False,
-            description=description,
-        )
-
-        # Verify debug message was logged
-        mock_logger.debug.assert_called_with(
-            "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format(
-                "test-statement-123", 1
-            )
-        )
-
-        # 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_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.download_manager = Mock()
-
-        # Mock the dependencies
-        mock_table = Mock()
-        queue._create_table_at_offset = 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_at_offset.assert_called_once_with(50)
-
-        # Verify progress was called
-        queue._progress_chunk_link.assert_called_once()
-
-        # Verify the result is the table
-        assert result == mock_table