fix type errors (ssl_options, CHUNK_PATH_WITH_ID..., etc.)

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

Author: varun-edachali-dbx

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

@@ -4,6 +4,7 @@
 import re
 from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING, Set
 
+from databricks.sql.backend.sea.models.base import ExternalLink
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
     ResultFormat,
@@ -91,6 +92,7 @@ class SeaDatabricksClient(DatabricksClient):
     STATEMENT_PATH = BASE_PATH + "statements"
     STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}"
     CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
+    CHUNK_PATH_WITH_ID_AND_INDEX = STATEMENT_PATH + "/{}/result/chunks/{}"
 
     def __init__(
         self,
@@ -326,7 +328,7 @@ def _extract_description_from_manifest(self, manifest_obj) -> Optional[List]:
 
         return columns if columns else None
 
-    def get_chunk_link(self, statement_id: str, chunk_index: int) -> "ExternalLink":
+    def get_chunk_link(self, statement_id: str, chunk_index: int) -> ExternalLink:
         """
         Get links for chunks starting from the specified index.
 
@@ -347,7 +349,13 @@ def get_chunk_link(self, statement_id: str, chunk_index: int) -> "ExternalLink":
         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}")
+            raise ServerOperationError(
+                f"No link found for chunk index {chunk_index}",
+                {
+                    "operation-id": statement_id,
+                    "diagnostic-info": None,
+                },
+            )
 
         return link
 

src/databricks/sql/result_set.py

@@ -488,6 +488,7 @@ def __init__(
                 result_data,
                 manifest,
                 str(execute_response.command_id.to_sea_statement_id()),
+                ssl_options=self.connection.session.ssl_options,
                 description=execute_response.description,
                 max_download_threads=sea_client.max_download_threads,
                 sea_client=sea_client,
@@ -512,27 +513,6 @@ def __init__(
         # Initialize queue for result data if not provided
         self.results = results_queue or JsonQueue([])
 
-    def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
-        """
-        Fetch the next set of rows as an Arrow table.
-
-        Args:
-            size: Number of rows to fetch
-
-        Returns:
-            PyArrow Table containing the fetched rows
-
-        Raises:
-            ImportError: If PyArrow is not installed
-            ValueError: If size is negative
-        """
-        if size < 0:
-            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
-
-        results = self.results.next_n_rows(size)
-        n_remaining_rows = size - results.num_rows
-        self._next_row_index += results.num_rows
-
     def fetchmany_json(self, size: int):
         """
         Fetch the next set of rows as a columnar table.

src/databricks/sql/session.py

@@ -64,7 +64,7 @@ def __init__(
         base_headers = [("User-Agent", useragent_header)]
         all_headers = (http_headers or []) + base_headers
 
-        self._ssl_options = SSLOptions(
+        self.ssl_options = SSLOptions(
             # Double negation is generally a bad thing, but we have to keep backward compatibility
             tls_verify=not kwargs.get(
                 "_tls_no_verify", False
@@ -113,7 +113,7 @@ def _create_backend(
             "http_path": http_path,
             "http_headers": all_headers,
             "auth_provider": auth_provider,
-            "ssl_options": self._ssl_options,
+            "ssl_options": self.ssl_options,
             "_use_arrow_native_complex_types": _use_arrow_native_complex_types,
             **kwargs,
         }

src/databricks/sql/utils.py

@@ -131,6 +131,7 @@ def build_queue(
         sea_result_data: ResultData,
         manifest: Optional[ResultManifest],
         statement_id: str,
+        ssl_options: Optional[SSLOptions] = None,
         description: Optional[List[Tuple[Any, ...]]] = None,
         max_download_threads: Optional[int] = None,
         sea_client: Optional["SeaDatabricksClient"] = None,