account for total chunk count

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

Author: varun-edachali-dbx

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

@@ -5,7 +5,7 @@
 """
 
 import base64
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Optional
 from dataclasses import dataclass
 
 from databricks.sql.backend.types import CommandState

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

@@ -159,6 +159,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(
@@ -169,8 +170,8 @@ def __init__(
         initial_links = result_data.external_links
         self._chunk_index_to_link = {link.chunk_index: link for link in initial_links}
 
-        initial_link = self._chunk_index_to_link.get(0, None)
-        if not initial_link:
+        first_link = self._chunk_index_to_link.get(0, None)
+        if not first_link:
             # possibly an empty response
             return
 
@@ -201,6 +202,11 @@ def _convert_to_thrift_link(self, link: "ExternalLink") -> TSparkArrowResultLink
         )
 
     def _get_chunk_link(self, chunk_index: int) -> Optional["ExternalLink"]:
+        if chunk_index >= self._total_chunk_count:
+            raise ValueError(
+                f"Chunk index {chunk_index} is greater than total chunk count {self._total_chunk_count}"
+            )
+
         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})