Merge branch 'sea-migration' into ext-links-sea

Author: varun-edachali-dbx

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

@@ -17,7 +17,7 @@
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
-    from databricks.sql.result_set import SeaResultSet
+    from databricks.sql.backend.sea.result_set import SeaResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.backend.types import (
@@ -253,7 +253,7 @@ def close_session(self, session_id: SessionId) -> None:
         logger.debug("SeaDatabricksClient.close_session(session_id=%s)", session_id)
 
         if session_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA session ID")
+            raise ValueError("Not a valid SEA session ID")
         sea_session_id = session_id.to_sea_session_id()
 
         request_data = DeleteSessionRequest(
@@ -292,7 +292,7 @@ def get_allowed_session_configurations() -> List[str]:
 
     def _extract_description_from_manifest(
         self, manifest: ResultManifest
-    ) -> Optional[List]:
+    ) -> List[Tuple]:
         """
         Extract column description from a manifest object, in the format defined by
         the spec: https://peps.python.org/pep-0249/#description
@@ -301,15 +301,12 @@ def _extract_description_from_manifest(
             manifest: The ResultManifest object containing schema information
 
         Returns:
-            Optional[List]: A list of column tuples or None if no columns are found
+            List[Tuple]: A list of column tuples
         """
 
         schema_data = manifest.schema
         columns_data = schema_data.get("columns", [])
 
-        if not columns_data:
-            return None
-
         columns = []
         for col_data in columns_data:
             # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
@@ -325,7 +322,7 @@ def _extract_description_from_manifest(
                 )
             )
 
-        return columns if columns else None
+        return columns
 
     def _results_message_to_execute_response(
         self, response: GetStatementResponse
@@ -431,7 +428,7 @@ def execute_command(
         """
 
         if session_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA session ID")
+            raise ValueError("Not a valid SEA session ID")
 
         sea_session_id = session_id.to_sea_session_id()
 
@@ -510,9 +507,11 @@ def cancel_command(self, command_id: CommandId) -> None:
         """
 
         if command_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA command ID")
+            raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
+        if sea_statement_id is None:
+            raise ValueError("Not a valid SEA command ID")
 
         request = CancelStatementRequest(statement_id=sea_statement_id)
         self.http_client._make_request(
@@ -533,9 +532,11 @@ def close_command(self, command_id: CommandId) -> None:
         """
 
         if command_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA command ID")
+            raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
+        if sea_statement_id is None:
+            raise ValueError("Not a valid SEA command ID")
 
         request = CloseStatementRequest(statement_id=sea_statement_id)
         self.http_client._make_request(
@@ -562,6 +563,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             raise ProgrammingError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
+        if sea_statement_id is None:
+            raise ValueError("Not a valid SEA command ID")
 
         request = GetStatementRequest(statement_id=sea_statement_id)
         response_data = self.http_client._make_request(
@@ -594,9 +597,11 @@ def get_execution_result(
         """
 
         if command_id.backend_type != BackendType.SEA:
-            raise ProgrammingError("Not a valid SEA command ID")
+            raise ValueError("Not a valid SEA command ID")
 
         sea_statement_id = command_id.to_sea_statement_id()
+        if sea_statement_id is None:
+            raise ValueError("Not a valid SEA command ID")
 
         # Create the request model
         request = GetStatementRequest(statement_id=sea_statement_id)
@@ -610,7 +615,7 @@ def get_execution_result(
         response = GetStatementResponse.from_dict(response_data)
 
         # Create and return a SeaResultSet
-        from databricks.sql.result_set import SeaResultSet
+        from databricks.sql.backend.sea.result_set import SeaResultSet
 
         execute_response = self._results_message_to_execute_response(response)
 

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

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import List, Optional, Tuple
+
+from databricks.sql.backend.sea.backend import SeaDatabricksClient
+from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+from databricks.sql.backend.sea.utils.constants import ResultFormat
+from databricks.sql.exc import ProgrammingError
+from databricks.sql.utils import ResultSetQueue
+
+
+class SeaResultSetQueueFactory(ABC):
+    @staticmethod
+    def build_queue(
+        sea_result_data: ResultData,
+        manifest: ResultManifest,
+        statement_id: str,
+        description: List[Tuple] = [],
+        max_download_threads: Optional[int] = None,
+        sea_client: Optional[SeaDatabricksClient] = None,
+        lz4_compressed: bool = False,
+    ) -> ResultSetQueue:
+        """
+        Factory method to build a result set queue for SEA backend.
+
+        Args:
+            sea_result_data (ResultData): Result data from SEA response
+            manifest (ResultManifest): Manifest from SEA response
+            statement_id (str): Statement ID for the query
+            description (List[List[Any]]): Column descriptions
+            max_download_threads (int): Maximum number of download threads
+            sea_client (SeaDatabricksClient): SEA client for fetching additional links
+            lz4_compressed (bool): Whether the data is LZ4 compressed
+
+        Returns:
+            ResultSetQueue: The appropriate queue for the result data
+        """
+
+        if manifest.format == ResultFormat.JSON_ARRAY.value:
+            # INLINE disposition with JSON_ARRAY format
+            return JsonQueue(sea_result_data.data)
+        elif manifest.format == ResultFormat.ARROW_STREAM.value:
+            # EXTERNAL_LINKS disposition
+            raise NotImplementedError(
+                "EXTERNAL_LINKS disposition is not implemented for SEA backend"
+            )
+        raise ProgrammingError("Invalid result format")
+
+
+class JsonQueue(ResultSetQueue):
+    """Queue implementation for JSON_ARRAY format data."""
+
+    def __init__(self, data_array: Optional[List[List[str]]]):
+        """Initialize with JSON array data."""
+        self.data_array = data_array or []
+        self.cur_row_index = 0
+        self.num_rows = len(self.data_array)
+
+    def next_n_rows(self, num_rows: int) -> List[List[str]]:
+        """Get the next n rows from the data array."""
+        length = min(num_rows, self.num_rows - self.cur_row_index)
+        slice = self.data_array[self.cur_row_index : self.cur_row_index + length]
+        self.cur_row_index += length
+        return slice
+
+    def remaining_rows(self) -> List[List[str]]:
+        """Get all remaining rows from the data array."""
+        slice = self.data_array[self.cur_row_index :]
+        self.cur_row_index += len(slice)
+        return slice