reduce responsibility of Queue

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

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

@@ -43,6 +43,7 @@
     ExecuteStatementResponse,
     GetStatementResponse,
     CreateSessionResponse,
+    GetChunksResponse,
 )
 
 logger = logging.getLogger(__name__)
@@ -305,9 +306,7 @@ def get_allowed_session_configurations() -> List[str]:
         """
         return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
 
-    def get_chunk_links(
-        self, statement_id: str, chunk_index: int
-    ) -> "GetChunksResponse":
+    def get_chunk_link(self, statement_id: str, chunk_index: int) -> "ExternalLink":
         """
         Get links for chunks starting from the specified index.
 
@@ -316,16 +315,21 @@ def get_chunk_links(
             chunk_index: The starting chunk index
 
         Returns:
-            GetChunksResponse: Response containing external links
+            ExternalLink: External link for the chunk
         """
-        from databricks.sql.backend.sea.models.responses import GetChunksResponse
 
         response_data = self.http_client._make_request(
             method="GET",
             path=self.CHUNK_PATH_WITH_ID_AND_INDEX.format(statement_id, chunk_index),
         )
+        response = GetChunksResponse.from_dict(response_data)
 
-        return GetChunksResponse.from_dict(response_data)
+        links = response.external_links
+        link = next((l for l in links if l.chunk_index == chunk_index), None)
+        if not link:
+            raise Error(f"No link found for chunk index {chunk_index}")
+
+        return link
 
     def _get_schema_bytes(self, sea_response) -> Optional[bytes]:
         """

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

@@ -27,6 +27,7 @@
     ExecuteStatementResponse,
     GetStatementResponse,
     CreateSessionResponse,
+    GetChunksResponse,
 )
 
 __all__ = [
@@ -49,4 +50,5 @@
     "ExecuteStatementResponse",
     "GetStatementResponse",
     "CreateSessionResponse",
+    "GetChunksResponse",
 ]

src/databricks/sql/cloud_fetch_queue.py

@@ -381,30 +381,19 @@ def _fetch_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
         )
 
         # Use the SEA client to fetch the chunk links
-        chunk_info = self._sea_client.get_chunk_links(self._statement_id, chunk_index)
-        links = chunk_info.external_links
+        link = self._sea_client.get_chunk_link(self._statement_id, chunk_index)
 
-        if not links:
-            logger.debug(
-                "SeaCloudFetchQueue: No links found for chunk {}".format(chunk_index)
-            )
-            return None
-
-        # Get the link for the requested chunk
-        link = next((l for l in links if l.chunk_index == chunk_index), None)
-
-        if link:
-            logger.debug(
-                "SeaCloudFetchQueue: Link details for chunk {}: row_offset={}, row_count={}, next_chunk_index={}".format(
-                    link.chunk_index,
-                    link.row_offset,
-                    link.row_count,
-                    link.next_chunk_index,
-                )
+        logger.debug(
+            "SeaCloudFetchQueue: Link details for chunk {}: row_offset={}, row_count={}, next_chunk_index={}".format(
+                link.chunk_index,
+                link.row_offset,
+                link.row_count,
+                link.next_chunk_index,
             )
+        )
 
-            if self.download_manager:
-                self.download_manager.add_links(self._convert_to_thrift_links([link]))
+        if self.download_manager:
+            self.download_manager.add_links(self._convert_to_thrift_links([link]))
 
         return link