Merge branch 'sea-migration' into metadata-sea

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -82,7 +82,7 @@ def execute_command(
         lz4_compression: bool,
         cursor: "Cursor",
         use_cloud_fetch: bool,
-        parameters: List[ttypes.TSparkParameter],
+        parameters: List,
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
     ) -> Union["ResultSet", None]:

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

@@ -1,8 +1,9 @@
 import logging
 import time
 import re
-from typing import Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
+from typing import Any, Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
 
+from databricks.sql.backend.sea.models.base import ResultManifest
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
     ResultFormat,
@@ -23,7 +24,7 @@
     BackendType,
     ExecuteResponse,
 )
-from databricks.sql.exc import ServerOperationError
+from databricks.sql.exc import DatabaseError, ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
 from databricks.sql.types import SSLOptions
 
@@ -89,6 +90,9 @@ class SeaDatabricksClient(DatabricksClient):
     STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}"
     CANCEL_STATEMENT_PATH_WITH_ID = STATEMENT_PATH + "/{}/cancel"
 
+    # SEA constants
+    POLL_INTERVAL_SECONDS = 0.2
+
     def __init__(
         self,
         server_hostname: str,
@@ -286,28 +290,28 @@ def get_allowed_session_configurations() -> List[str]:
         """
         return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
 
-    def _extract_description_from_manifest(self, manifest_obj) -> Optional[List]:
+    def _extract_description_from_manifest(
+        self, manifest: ResultManifest
+    ) -> Optional[List]:
         """
-        Extract column description from a manifest object.
+        Extract column description from a manifest object, in the format defined by
+        the spec: https://peps.python.org/pep-0249/#description
 
         Args:
-            manifest_obj: The ResultManifest object containing schema information
+            manifest: The ResultManifest object containing schema information
 
         Returns:
             Optional[List]: A list of column tuples or None if no columns are found
         """
 
-        schema_data = manifest_obj.schema
+        schema_data = manifest.schema
         columns_data = schema_data.get("columns", [])
 
         if not columns_data:
             return None
 
         columns = []
         for col_data in columns_data:
-            if not isinstance(col_data, dict):
-                continue
-
             # Format: (name, type_code, display_size, internal_size, precision, scale, null_ok)
             columns.append(
                 (
@@ -323,7 +327,9 @@ def _extract_description_from_manifest(self, manifest_obj) -> Optional[List]:
 
         return columns if columns else None
 
-    def _results_message_to_execute_response(self, sea_response, command_id):
+    def _results_message_to_execute_response(
+        self, response: GetStatementResponse
+    ) -> ExecuteResponse:
         """
         Convert a SEA response to an ExecuteResponse and extract result data.
 
@@ -332,33 +338,65 @@ def _results_message_to_execute_response(self, sea_response, command_id):
             command_id: The command ID
 
         Returns:
-            tuple: (ExecuteResponse, ResultData, ResultManifest) - The normalized execute response,
-                  result data object, and manifest object
+            ExecuteResponse: The normalized execute response
         """
 
-        # Parse the response
-        status = parse_status(sea_response)
-        manifest_obj = parse_manifest(sea_response)
-        result_data_obj = parse_result(sea_response)
-
         # Extract description from manifest schema
-        description = self._extract_description_from_manifest(manifest_obj)
+        description = self._extract_description_from_manifest(response.manifest)
 
         # Check for compression
-        lz4_compressed = manifest_obj.result_compression == "LZ4_FRAME"
+        lz4_compressed = (
+            response.manifest.result_compression == ResultCompression.LZ4_FRAME
+        )
 
         execute_response = ExecuteResponse(
-            command_id=command_id,
-            status=status.state,
+            command_id=CommandId.from_sea_statement_id(response.statement_id),
+            status=response.status.state,
             description=description,
             has_been_closed_server_side=False,
             lz4_compressed=lz4_compressed,
             is_staging_operation=False,
             arrow_schema_bytes=None,
-            result_format=manifest_obj.format,
+            result_format=response.manifest.format,
         )
 
-        return execute_response, result_data_obj, manifest_obj
+        return execute_response
+
+    def _check_command_not_in_failed_or_closed_state(
+        self, state: CommandState, command_id: CommandId
+    ) -> None:
+        if state == CommandState.CLOSED:
+            raise DatabaseError(
+                "Command {} unexpectedly closed server side".format(command_id),
+                {
+                    "operation-id": command_id,
+                },
+            )
+        if state == CommandState.FAILED:
+            raise ServerOperationError(
+                "Command {} failed".format(command_id),
+                {
+                    "operation-id": command_id,
+                },
+            )
+
+    def _wait_until_command_done(
+        self, response: ExecuteStatementResponse
+    ) -> CommandState:
+        """
+        Wait until a command is done.
+        """
+
+        state = response.status.state
+        command_id = CommandId.from_sea_statement_id(response.statement_id)
+
+        while state in [CommandState.PENDING, CommandState.RUNNING]:
+            time.sleep(self.POLL_INTERVAL_SECONDS)
+            state = self.get_query_state(command_id)
+
+        self._check_command_not_in_failed_or_closed_state(state, command_id)
+
+        return state
 
     def execute_command(
         self,
@@ -369,7 +407,7 @@ def execute_command(
         lz4_compression: bool,
         cursor: "Cursor",
         use_cloud_fetch: bool,
-        parameters: List,
+        parameters: List[Dict[str, Any]],
         async_op: bool,
         enforce_embedded_schema_correctness: bool,
     ) -> Union["ResultSet", None]:
@@ -403,9 +441,9 @@ def execute_command(
             for param in parameters:
                 sea_parameters.append(
                     StatementParameter(
-                        name=param.name,
-                        value=param.value,
-                        type=param.type if hasattr(param, "type") else None,
+                        name=param["name"],
+                        value=param["value"],
+                        type=param["type"] if "type" in param else None,
                     )
                 )
 
@@ -457,24 +495,7 @@ def execute_command(
         if async_op:
             return None
 
-        # For synchronous operation, wait for the statement to complete
-        status = response.status
-        state = status.state
-
-        # Keep polling until we reach a terminal state
-        while state in [CommandState.PENDING, CommandState.RUNNING]:
-            time.sleep(0.5)  # add a small delay to avoid excessive API calls
-            state = self.get_query_state(command_id)
-
-        if state != CommandState.SUCCEEDED:
-            raise ServerOperationError(
-                f"Statement execution did not succeed: {status.error.message if status.error else 'Unknown error'}",
-                {
-                    "operation-id": command_id.to_sea_statement_id(),
-                    "diagnostic-info": None,
-                },
-            )
-
+        self._wait_until_command_done(response)
         return self.get_execution_result(command_id, cursor)
 
     def cancel_command(self, command_id: CommandId) -> None:
@@ -586,25 +607,21 @@ def get_execution_result(
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
         )
+        response = GetStatementResponse.from_dict(response_data)
 
         # Create and return a SeaResultSet
         from databricks.sql.result_set import SeaResultSet
 
-        # Convert the response to an ExecuteResponse and extract result data
-        (
-            execute_response,
-            result_data,
-            manifest,
-        ) = self._results_message_to_execute_response(response_data, command_id)
+        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=result_data,
-            manifest=manifest,
+            result_data=response.result,
+            manifest=response.manifest,
         )
 
     # == Metadata Operations ==

tests/unit/test_sea_backend.py

@@ -15,7 +15,12 @@
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
 from databricks.sql.types import SSLOptions
 from databricks.sql.auth.authenticators import AuthProvider
-from databricks.sql.exc import Error, NotSupportedError, ServerOperationError
+from databricks.sql.exc import (
+    Error,
+    NotSupportedError,
+    ServerOperationError,
+    DatabaseError,
+)
 
 
 class TestSeaBackend:
@@ -349,10 +354,7 @@ def test_command_execution_advanced(
             "status": {"state": "SUCCEEDED"},
         }
         mock_http_client._make_request.return_value = execute_response
-        param = MagicMock()
-        param.name = "param1"
-        param.value = "value1"
-        param.type = "STRING"
+        param = {"name": "param1", "value": "value1", "type": "STRING"}
 
         with patch.object(sea_client, "get_execution_result"):
             sea_client.execute_command(
@@ -405,7 +407,7 @@ def test_command_execution_advanced(
                         async_op=False,
                         enforce_embedded_schema_correctness=False,
                     )
-                assert "Statement execution did not succeed" in str(excinfo.value)
+                assert "Command test-statement-123 failed" in str(excinfo.value)
 
         # Test missing statement ID
         mock_http_client.reset_mock()
@@ -523,6 +525,34 @@ def test_command_management(
             sea_client.get_execution_result(thrift_command_id, mock_cursor)
         assert "Not a valid SEA command ID" in str(excinfo.value)
 
+    def test_check_command_state(self, sea_client, sea_command_id):
+        """Test _check_command_not_in_failed_or_closed_state method."""
+        # Test with RUNNING state (should not raise)
+        sea_client._check_command_not_in_failed_or_closed_state(
+            CommandState.RUNNING, sea_command_id
+        )
+
+        # Test with SUCCEEDED state (should not raise)
+        sea_client._check_command_not_in_failed_or_closed_state(
+            CommandState.SUCCEEDED, sea_command_id
+        )
+
+        # Test with CLOSED state (should raise DatabaseError)
+        with pytest.raises(DatabaseError) as excinfo:
+            sea_client._check_command_not_in_failed_or_closed_state(
+                CommandState.CLOSED, sea_command_id
+            )
+        assert "Command test-statement-123 unexpectedly closed server side" in str(
+            excinfo.value
+        )
+
+        # Test with FAILED state (should raise ServerOperationError)
+        with pytest.raises(ServerOperationError) as excinfo:
+            sea_client._check_command_not_in_failed_or_closed_state(
+                CommandState.FAILED, sea_command_id
+            )
+        assert "Command test-statement-123 failed" in str(excinfo.value)
+
     def test_utility_methods(self, sea_client):
         """Test utility methods."""
         # Test get_default_session_configuration_value
@@ -590,12 +620,66 @@ def test_utility_methods(self, sea_client):
         assert description[1][1] == "INT"  # type_code
         assert description[1][6] is False  # null_ok
 
-        # Test with manifest containing non-dict column
-        manifest_obj.schema = {"columns": ["not_a_dict"]}
-        description = sea_client._extract_description_from_manifest(manifest_obj)
-        assert description is None
+        # Test _extract_description_from_manifest with empty columns
+        empty_manifest = MagicMock()
+        empty_manifest.schema = {"columns": []}
+        assert sea_client._extract_description_from_manifest(empty_manifest) is None
 
-        # Test with manifest without columns
-        manifest_obj.schema = {}
-        description = sea_client._extract_description_from_manifest(manifest_obj)
-        assert description is None
+        # Test _extract_description_from_manifest with no columns key
+        no_columns_manifest = MagicMock()
+        no_columns_manifest.schema = {}
+        assert (
+            sea_client._extract_description_from_manifest(no_columns_manifest) is None
+        )
+
+    def test_unimplemented_metadata_methods(
+        self, sea_client, sea_session_id, mock_cursor
+    ):
+        """Test that metadata methods raise NotImplementedError."""
+        # Test get_catalogs
+        with pytest.raises(NotImplementedError):
+            sea_client.get_catalogs(sea_session_id, 100, 1000, mock_cursor)
+
+        # Test get_schemas
+        with pytest.raises(NotImplementedError):
+            sea_client.get_schemas(sea_session_id, 100, 1000, mock_cursor)
+
+        # Test get_schemas with optional parameters
+        with pytest.raises(NotImplementedError):
+            sea_client.get_schemas(
+                sea_session_id, 100, 1000, mock_cursor, "catalog", "schema"
+            )
+
+        # Test get_tables
+        with pytest.raises(NotImplementedError):
+            sea_client.get_tables(sea_session_id, 100, 1000, mock_cursor)
+
+        # Test get_tables with optional parameters
+        with pytest.raises(NotImplementedError):
+            sea_client.get_tables(
+                sea_session_id,
+                100,
+                1000,
+                mock_cursor,
+                catalog_name="catalog",
+                schema_name="schema",
+                table_name="table",
+                table_types=["TABLE", "VIEW"],
+            )
+
+        # Test get_columns
+        with pytest.raises(NotImplementedError):
+            sea_client.get_columns(sea_session_id, 100, 1000, mock_cursor)
+
+        # Test get_columns with optional parameters
+        with pytest.raises(NotImplementedError):
+            sea_client.get_columns(
+                sea_session_id,
+                100,
+                1000,
+                mock_cursor,
+                catalog_name="catalog",
+                schema_name="schema",
+                table_name="table",
+                column_name="column",
+            )