Merge branch 'sea-migration' into sea-sess-conf-fix

Author: varun-edachali-dbx

examples/experimental/tests/test_sea_async_query.py

@@ -51,12 +51,20 @@ def test_sea_async_query_with_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query asynchronously
+        # Execute a query that generates large rows to force multiple chunks
+        requested_row_count = 5000
         cursor = connection.cursor()
+        query = f"""
+        SELECT 
+            id, 
+            concat('value_', repeat('a', 10000)) as test_value
+        FROM range(1, {requested_row_count} + 1) AS t(id)
+        """
+
         logger.info(
-            "Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
+            f"Executing asynchronous query with cloud fetch to generate {requested_row_count} rows"
         )
-        cursor.execute_async("SELECT 1 as test_value")
+        cursor.execute_async(query)
         logger.info(
             "Asynchronous query submitted successfully with cloud fetch enabled"
         )
@@ -69,8 +77,25 @@ def test_sea_async_query_with_cloud_fetch():
 
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
+
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        actual_row_count = len(results)
+
+        logger.info(
+            f"Requested {requested_row_count} rows, received {actual_row_count} rows"
+        )
+
+        # Verify total row count
+        if actual_row_count != requested_row_count:
+            logger.error(
+                f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
+            return False
+
         logger.info(
-            "Successfully retrieved asynchronous query results with cloud fetch enabled"
+            "PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly"
         )
 
         # Close resources
@@ -130,12 +155,20 @@ def test_sea_async_query_without_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query asynchronously
+        # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
+        requested_row_count = 100
         cursor = connection.cursor()
+        query = f"""
+        SELECT 
+            id, 
+            concat('value_', repeat('a', 100)) as test_value
+        FROM range(1, {requested_row_count} + 1) AS t(id)
+        """
+
         logger.info(
-            "Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
+            f"Executing asynchronous query without cloud fetch to generate {requested_row_count} rows"
         )
-        cursor.execute_async("SELECT 1 as test_value")
+        cursor.execute_async(query)
         logger.info(
             "Asynchronous query submitted successfully with cloud fetch disabled"
         )
@@ -148,8 +181,24 @@ def test_sea_async_query_without_cloud_fetch():
 
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        actual_row_count = len(results)
+
+        logger.info(
+            f"Requested {requested_row_count} rows, received {actual_row_count} rows"
+        )
+
+        # Verify total row count
+        if actual_row_count != requested_row_count:
+            logger.error(
+                f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
+            return False
+
         logger.info(
-            "Successfully retrieved asynchronous query results with cloud fetch disabled"
+            "PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly"
         )
 
         # Close resources

examples/experimental/tests/test_sea_sync_query.py

@@ -49,13 +49,27 @@ def test_sea_sync_query_with_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query
+        # Execute a query that generates large rows to force multiple chunks
+        requested_row_count = 10000
         cursor = connection.cursor()
+        query = f"""
+        SELECT 
+            id, 
+            concat('value_', repeat('a', 10000)) as test_value
+        FROM range(1, {requested_row_count} + 1) AS t(id)
+        """
+
+        logger.info(
+            f"Executing synchronous query with cloud fetch to generate {requested_row_count} rows"
+        )
+        cursor.execute(query)
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        actual_row_count = len(results)
         logger.info(
-            "Executing synchronous query with cloud fetch: SELECT 1 as test_value"
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
         )
-        cursor.execute("SELECT 1 as test_value")
-        logger.info("Query executed successfully with cloud fetch enabled")
 
         # Close resources
         cursor.close()
@@ -114,13 +128,18 @@ def test_sea_sync_query_without_cloud_fetch():
             f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
         )
 
-        # Execute a simple query
+        # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
+        requested_row_count = 100
         cursor = connection.cursor()
-        logger.info(
-            "Executing synchronous query without cloud fetch: SELECT 1 as test_value"
+        logger.info("Executing synchronous query without cloud fetch: SELECT 100 rows")
+        cursor.execute(
+            "SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
         )
-        cursor.execute("SELECT 1 as test_value")
-        logger.info("Query executed successfully with cloud fetch disabled")
+
+        results = [cursor.fetchone()]
+        results.extend(cursor.fetchmany(10))
+        results.extend(cursor.fetchall())
+        logger.info(f"{len(results)} rows retrieved against 100 requested")
 
         # Close resources
         cursor.close()

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 (
@@ -251,7 +251,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(
@@ -290,7 +290,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
@@ -299,15 +299,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)
@@ -323,7 +320,7 @@ def _extract_description_from_manifest(
                 )
             )
 
-        return columns if columns else None
+        return columns
 
     def _results_message_to_execute_response(
         self, response: GetStatementResponse
@@ -429,7 +426,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()
 
@@ -508,9 +505,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(
@@ -531,9 +530,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(
@@ -560,6 +561,8 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             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 = GetStatementRequest(statement_id=sea_statement_id)
         response_data = self.http_client._make_request(
@@ -592,9 +595,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)
@@ -608,18 +613,18 @@ 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)
 
         return SeaResultSet(
             connection=cursor.connection,
             execute_response=execute_response,
             sea_client=self,
-            buffer_size_bytes=cursor.buffer_size_bytes,
-            arraysize=cursor.arraysize,
             result_data=response.result,
             manifest=response.manifest,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
         )
 
     # == Metadata Operations ==

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