Merge branch 'sea-migration' into comparator

Author: varun-edachali-dbx

examples/experimental/tests/test_sea_sync_query.py

@@ -72,6 +72,13 @@ def test_sea_sync_query_with_cloud_fetch():
             f"{actual_row_count} rows retrieved against {requested_row_count} requested"
         )
 
+        # 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
+
         # Close resources
         cursor.close()
         connection.close()
@@ -132,15 +139,27 @@ def test_sea_sync_query_without_cloud_fetch():
         # 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 100 rows")
+        logger.info(
+            f"Executing synchronous query without cloud fetch: SELECT {requested_row_count} rows"
+        )
         cursor.execute(
             "SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
         )
 
         results = [cursor.fetchone()]
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
-        logger.info(f"{len(results)} rows retrieved against 100 requested")
+        actual_row_count = len(results)
+        logger.info(
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
+        )
+
+        # 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
 
         # Close resources
         cursor.close()

src/databricks/sql/backend/databricks_client.py

@@ -96,7 +96,7 @@ def execute_command(
             max_rows: Maximum number of rows to fetch in a single fetch batch
             max_bytes: Maximum number of bytes to fetch in a single fetch batch
             lz4_compression: Whether to use LZ4 compression for result data
-            cursor: The cursor object that will handle the results
+            cursor: The cursor object that will handle the results. The command id is set in this cursor.
             use_cloud_fetch: Whether to use cloud fetch for retrieving large result sets
             parameters: List of parameters to bind to the query
             async_op: Whether to execute the command asynchronously
@@ -282,7 +282,9 @@ def get_tables(
             max_bytes: Maximum number of bytes to fetch in a single batch
             cursor: The cursor object that will handle the results
             catalog_name: Optional catalog name pattern to filter by
+                if catalog_name is None, we fetch across all catalogs
             schema_name: Optional schema name pattern to filter by
+                if schema_name is None, we fetch across all schemas
             table_name: Optional table name pattern to filter by
             table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])
 
@@ -321,6 +323,7 @@ def get_columns(
             catalog_name: Optional catalog name pattern to filter by
             schema_name: Optional schema name pattern to filter by
             table_name: Optional table name pattern to filter by
+                if table_name is None, we fetch across all tables
             column_name: Optional column name pattern to filter by
 
         Returns:

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

@@ -5,7 +5,12 @@
 import re
 from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
 
-from databricks.sql.backend.sea.models.base import ExternalLink, ResultManifest
+from databricks.sql.backend.sea.models.base import (
+    ExternalLink,
+    ResultManifest,
+    StatementStatus,
+)
+from databricks.sql.backend.sea.models.responses import GetChunksResponse
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
     ResultFormat,
@@ -45,14 +50,34 @@
     GetStatementResponse,
     CreateSessionResponse,
 )
-from databricks.sql.backend.sea.models.responses import GetChunksResponse
 
 logger = logging.getLogger(__name__)
 
 
 def _filter_session_configuration(
     session_configuration: Optional[Dict[str, Any]],
 ) -> Dict[str, str]:
+    """
+    Filter and normalise the provided session configuration parameters.
+
+    The Statement Execution API supports only a subset of SQL session
+    configuration options.  This helper validates the supplied
+    ``session_configuration`` dictionary against the allow-list defined in
+    ``ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP`` and returns a new
+    dictionary that contains **only** the supported parameters.
+
+    Args:
+        session_configuration: Optional mapping of session configuration
+            names to their desired values.  Key comparison is
+            case-insensitive.
+
+    Returns:
+        Dict[str, str]: A dictionary containing only the supported
+        configuration parameters with lower-case keys and string values.  If
+        *session_configuration* is ``None`` or empty, an empty dictionary is
+        returned.
+    """
+
     if not session_configuration:
         return {}
 
@@ -143,7 +168,7 @@ def __init__(
             http_path=http_path,
             http_headers=http_headers,
             auth_provider=auth_provider,
-            ssl_options=self._ssl_options,
+            ssl_options=ssl_options,
             **kwargs,
         )
 
@@ -275,29 +300,6 @@ def close_session(self, session_id: SessionId) -> None:
             data=request_data.to_dict(),
         )
 
-    @staticmethod
-    def get_default_session_configuration_value(name: str) -> Optional[str]:
-        """
-        Get the default value for a session configuration parameter.
-
-        Args:
-            name: The name of the session configuration parameter
-
-        Returns:
-            The default value if the parameter is supported, None otherwise
-        """
-        return ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.get(name.upper())
-
-    @staticmethod
-    def get_allowed_session_configurations() -> List[str]:
-        """
-        Get the list of allowed session configuration parameters.
-
-        Returns:
-            List of allowed session configuration parameter names
-        """
-        return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
-
     def _extract_description_from_manifest(
         self, manifest: ResultManifest
     ) -> List[Tuple]:
@@ -309,7 +311,7 @@ def _extract_description_from_manifest(
             manifest: The ResultManifest object containing schema information
 
         Returns:
-            List[Tuple]: A list of column tuples
+            Optional[List]: A list of column tuples or None if no columns are found
         """
 
         schema_data = manifest.schema
@@ -397,8 +399,9 @@ def _response_to_result_set(
         )
 
     def _check_command_not_in_failed_or_closed_state(
-        self, state: CommandState, command_id: CommandId
+        self, status: StatementStatus, command_id: CommandId
     ) -> None:
+        state = status.state
         if state == CommandState.CLOSED:
             raise DatabaseError(
                 "Command {} unexpectedly closed server side".format(command_id),
@@ -407,8 +410,11 @@ def _check_command_not_in_failed_or_closed_state(
                 },
             )
         if state == CommandState.FAILED:
+            error = status.error
+            error_code = error.error_code if error else "UNKNOWN_ERROR_CODE"
+            error_message = error.message if error else "UNKNOWN_ERROR_MESSAGE"
             raise ServerOperationError(
-                "Command {} failed".format(command_id),
+                "Command failed: {} - {}".format(error_code, error_message),
                 {
                     "operation-id": command_id,
                 },
@@ -422,16 +428,18 @@ def _wait_until_command_done(
         """
 
         final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response
-
-        state = final_response.status.state
         command_id = CommandId.from_sea_statement_id(final_response.statement_id)
 
-        while state in [CommandState.PENDING, CommandState.RUNNING]:
+        while final_response.status.state in [
+            CommandState.PENDING,
+            CommandState.RUNNING,
+        ]:
             time.sleep(self.POLL_INTERVAL_SECONDS)
             final_response = self._poll_query(command_id)
-            state = final_response.status.state
 
-        self._check_command_not_in_failed_or_closed_state(state, command_id)
+        self._check_command_not_in_failed_or_closed_state(
+            final_response.status, command_id
+        )
 
         return final_response
 
@@ -465,7 +473,7 @@ def execute_command(
             enforce_embedded_schema_correctness: Whether to enforce schema correctness
 
         Returns:
-            SeaResultSet: A SeaResultSet instance for the executed command
+            ResultSet: A SeaResultSet instance for the executed command
         """
 
         if session_id.backend_type != BackendType.SEA:
@@ -521,14 +529,6 @@ def execute_command(
         )
         response = ExecuteStatementResponse.from_dict(response_data)
         statement_id = response.statement_id
-        if not statement_id:
-            raise ServerOperationError(
-                "Failed to execute command: No statement ID returned",
-                {
-                    "operation-id": None,
-                    "diagnostic-info": None,
-                },
-            )
 
         command_id = CommandId.from_sea_statement_id(statement_id)
 
@@ -560,8 +560,6 @@ def cancel_command(self, command_id: CommandId) -> None:
             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(
@@ -585,8 +583,6 @@ def close_command(self, command_id: CommandId) -> None:
             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(
@@ -604,8 +600,6 @@ def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
             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(
@@ -628,7 +622,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
             CommandState: The current state of the command
 
         Raises:
-            ProgrammingError: If the command ID is invalid
+            ValueError: If the command ID is invalid
         """
 
         response = self._poll_query(command_id)

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

@@ -9,7 +9,6 @@
     StatementStatus,
     ExternalLink,
     ResultData,
-    ColumnInfo,
     ResultManifest,
 )
 
@@ -36,7 +35,6 @@
     "StatementStatus",
     "ExternalLink",
     "ResultData",
-    "ColumnInfo",
     "ResultManifest",
     # Request models
     "StatementParameter",

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

@@ -67,25 +67,12 @@ class ResultData:
     attachment: Optional[bytes] = None
 
 
-@dataclass
-class ColumnInfo:
-    """Information about a column in the result set."""
-
-    name: str
-    type_name: str
-    type_text: str
-    nullable: bool = True
-    precision: Optional[int] = None
-    scale: Optional[int] = None
-    ordinal_position: Optional[int] = None
-
-
 @dataclass
 class ResultManifest:
     """Manifest information for a result set."""
 
     format: str
-    schema: Dict[str, Any]  # Will contain column information
+    schema: Dict[str, Any]
     total_row_count: int
     total_byte_count: int
     total_chunk_count: int

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

@@ -54,8 +54,8 @@ def to_dict(self) -> Dict[str, Any]:
             result["parameters"] = [
                 {
                     "name": param.name,
-                    **({"value": param.value} if param.value is not None else {}),
-                    **({"type": param.type} if param.type is not None else {}),
+                    "value": param.value,
+                    "type": param.type,
                 }
                 for param in self.parameters
             ]

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

@@ -72,6 +72,7 @@ def build_queue(
             return JsonQueue(result_data.data)
         elif manifest.format == ResultFormat.ARROW_STREAM.value:
             if result_data.attachment is not None:
+                # direct results from Hybrid disposition
                 arrow_file = (
                     ResultSetDownloadHandler._decompress_data(result_data.attachment)
                     if lz4_compressed

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

@@ -53,7 +53,7 @@ def _filter_sea_result_set(
         # Reuse the command_id from the original result set
         command_id = result_set.command_id
 
-        # Create an ExecuteResponse with the filtered data
+        # Create an ExecuteResponse for the filtered data
         execute_response = ExecuteResponse(
             command_id=command_id,
             status=result_set.status,