remove un-necessary Optional params

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

Author: varun-edachali-dbx

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

@@ -35,11 +35,11 @@ def build_queue(
         result_data: ResultData,
         manifest: ResultManifest,
         statement_id: str,
-        ssl_options: Optional[SSLOptions] = None,
-        description: List[Tuple] = [],
-        max_download_threads: Optional[int] = None,
-        sea_client: Optional[SeaDatabricksClient] = None,
-        lz4_compressed: bool = False,
+        ssl_options: SSLOptions,
+        description: List[Tuple],
+        max_download_threads: int,
+        sea_client: SeaDatabricksClient,
+        lz4_compressed: bool,
     ) -> ResultSetQueue:
         """
         Factory method to build a result set queue for SEA backend.
@@ -62,19 +62,6 @@ def build_queue(
             return JsonQueue(result_data.data)
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
             # EXTERNAL_LINKS disposition
-            if not max_download_threads:
-                raise ValueError(
-                    "Max download threads is required for EXTERNAL_LINKS disposition"
-                )
-            if not ssl_options:
-                raise ValueError(
-                    "SSL options are required for EXTERNAL_LINKS disposition"
-                )
-            if not sea_client:
-                raise ValueError(
-                    "SEA client is required for EXTERNAL_LINKS disposition"
-                )
-
             return SeaCloudFetchQueue(
                 initial_links=result_data.external_links or [],
                 max_download_threads=max_download_threads,

src/databricks/sql/utils.py

@@ -261,7 +261,6 @@ def next_n_rows(self, num_rows: int) -> "pyarrow.Table":
             length = min(num_rows, self.table.num_rows - self.table_row_index)
             table_slice = self.table.slice(self.table_row_index, length)
             results = pyarrow.concat_tables([results, table_slice])
-
             self.table_row_index += table_slice.num_rows
 
             # Replace current table with the next table if we are at the end of the current table
@@ -284,21 +283,19 @@ def remaining_rows(self) -> "pyarrow.Table":
         if not self.table:
             # Return empty pyarrow table to cause retry of fetch
             return self._create_empty_table()
-
         results = self.table.slice(0, 0)
         while self.table:
             table_slice = self.table.slice(
                 self.table_row_index, self.table.num_rows - self.table_row_index
             )
             results = pyarrow.concat_tables([results, table_slice])
-
             self.table_row_index += table_slice.num_rows
             self.table = self._create_next_table()
             self.table_row_index = 0
         return results
 
     def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
-        """Create next table by retrieving the logical next downloaded file."""
+        """Create next table at the given row offset"""
         # Create next table by retrieving the logical next downloaded file, or return None to signal end of queue
         if not self.download_manager:
             logger.debug("CloudFetchQueue: No download manager available")
@@ -311,7 +308,6 @@ def _create_table_at_offset(self, offset: int) -> Union["pyarrow.Table", None]:
             )
             # None signals no more Arrow tables can be built from the remaining handlers if any remain
             return None
-
         arrow_table = create_arrow_table_from_arrow_file(
             downloaded_file.file_bytes, self.description
         )

tests/unit/test_sea_queue.py

@@ -183,6 +183,11 @@ def test_build_queue_json_array(self, json_manifest, sample_data):
             result_data=result_data,
             manifest=json_manifest,
             statement_id="test-statement",
+            ssl_options=SSLOptions(),
+            description=[],
+            max_download_threads=10,
+            sea_client=Mock(),
+            lz4_compressed=False,
         )
 
         assert isinstance(queue, JsonQueue)
@@ -222,49 +227,6 @@ def test_build_queue_arrow_stream(
 
         assert isinstance(queue, SeaCloudFetchQueue)
 
-    def test_build_queue_arrow_stream_missing_threads(
-        self, arrow_manifest, ssl_options, mock_sea_client
-    ):
-        """Test building an Arrow stream queue with missing max_download_threads."""
-        result_data = ResultData(data=None, external_links=[])
-
-        with pytest.raises(ValueError, match="Max download threads is required"):
-            SeaResultSetQueueFactory.build_queue(
-                result_data=result_data,
-                manifest=arrow_manifest,
-                statement_id="test-statement",
-                ssl_options=ssl_options,
-                sea_client=mock_sea_client,
-            )
-
-    def test_build_queue_arrow_stream_missing_ssl(
-        self, arrow_manifest, mock_sea_client
-    ):
-        """Test building an Arrow stream queue with missing SSL options."""
-        result_data = ResultData(data=None, external_links=[])
-
-        with pytest.raises(ValueError, match="SSL options are required"):
-            SeaResultSetQueueFactory.build_queue(
-                result_data=result_data,
-                manifest=arrow_manifest,
-                statement_id="test-statement",
-                max_download_threads=10,
-                sea_client=mock_sea_client,
-            )
-
-    def test_build_queue_arrow_stream_missing_client(self, arrow_manifest, ssl_options):
-        """Test building an Arrow stream queue with missing SEA client."""
-        result_data = ResultData(data=None, external_links=[])
-
-        with pytest.raises(ValueError, match="SEA client is required"):
-            SeaResultSetQueueFactory.build_queue(
-                result_data=result_data,
-                manifest=arrow_manifest,
-                statement_id="test-statement",
-                ssl_options=ssl_options,
-                max_download_threads=10,
-            )
-
     def test_build_queue_invalid_format(self, invalid_manifest):
         """Test building a queue with invalid format."""
         result_data = ResultData(data=[])
@@ -274,6 +236,11 @@ def test_build_queue_invalid_format(self, invalid_manifest):
                 result_data=result_data,
                 manifest=invalid_manifest,
                 statement_id="test-statement",
+                ssl_options=SSLOptions(),
+                description=[],
+                max_download_threads=10,
+                sea_client=Mock(),
+                lz4_compressed=False,
             )