Implement SeaDatabricksClient (Complete Execution Spec)

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

@@ -1,25 +1,45 @@
 import logging
+import time
 import re
-from typing import Dict, Tuple, List, Optional, TYPE_CHECKING, Set
+from typing import Dict, Tuple, List, Optional, Union, TYPE_CHECKING, Set
+
+from databricks.sql.backend.sea.utils.constants import (
+    ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
+)
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
+    from databricks.sql.result_set import ResultSet
 
 from databricks.sql.backend.databricks_client import DatabricksClient
-from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
+from databricks.sql.backend.types import (
+    SessionId,
+    CommandId,
+    CommandState,
+    BackendType,
+    ExecuteResponse,
+)
 from databricks.sql.exc import ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
-from databricks.sql.backend.sea.utils.constants import (
-    ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
-)
-from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
 
 from databricks.sql.backend.sea.models import (
+    ExecuteStatementRequest,
+    GetStatementRequest,
+    CancelStatementRequest,
+    CloseStatementRequest,
     CreateSessionRequest,
     DeleteSessionRequest,
+    StatementParameter,
+    ExecuteStatementResponse,
+    GetStatementResponse,
     CreateSessionResponse,
 )
+from databricks.sql.backend.sea.models.responses import (
+    parse_status,
+    parse_manifest,
+    parse_result,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -262,8 +282,79 @@ def get_allowed_session_configurations() -> List[str]:
         """
         return list(ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP.keys())
 
-    # == Not Implemented Operations ==
-    # These methods will be implemented in future iterations
+    def _extract_description_from_manifest(self, manifest_obj) -> Optional[List]:
+        """
+        Extract column description from a manifest object.
+
+        Args:
+            manifest_obj: 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
+        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(
+                (
+                    col_data.get("name", ""),  # name
+                    col_data.get("type_name", ""),  # type_code
+                    None,  # display_size (not provided by SEA)
+                    None,  # internal_size (not provided by SEA)
+                    col_data.get("precision"),  # precision
+                    col_data.get("scale"),  # scale
+                    col_data.get("nullable", True),  # null_ok
+                )
+            )
+
+        return columns if columns else None
+
+    def _results_message_to_execute_response(self, sea_response, command_id):
+        """
+        Convert a SEA response to an ExecuteResponse and extract result data.
+
+        Args:
+            sea_response: The response from the SEA API
+            command_id: The command ID
+
+        Returns:
+            tuple: (ExecuteResponse, ResultData, ResultManifest) - The normalized execute response,
+                  result data object, and manifest object
+        """
+
+        # 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)
+
+        # Check for compression
+        lz4_compressed = manifest_obj.result_compression == "LZ4_FRAME"
+
+        execute_response = ExecuteResponse(
+            command_id=command_id,
+            status=status.state,
+            description=description,
+            has_been_closed_server_side=False,
+            lz4_compressed=lz4_compressed,
+            is_staging_operation=False,
+            arrow_schema_bytes=None,  # to be extracted during fetch phase for ARROW
+            result_format=manifest_obj.format,
+        )
+
+        return execute_response, result_data_obj, manifest_obj
 
     def execute_command(
         self,
@@ -274,41 +365,234 @@ 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,
-    ):
-        """Not implemented yet."""
-        raise NotImplementedError(
-            "execute_command is not yet implemented for SEA backend"
+    ) -> Union["ResultSet", None]:
+        """
+        Execute a SQL command using the SEA backend.
+
+        Args:
+            operation: SQL command to execute
+            session_id: Session identifier
+            max_rows: Maximum number of rows to fetch
+            max_bytes: Maximum number of bytes to fetch
+            lz4_compression: Whether to use LZ4 compression
+            cursor: Cursor executing the command
+            use_cloud_fetch: Whether to use cloud fetch
+            parameters: SQL parameters
+            async_op: Whether to execute asynchronously
+            enforce_embedded_schema_correctness: Whether to enforce schema correctness
+
+        Returns:
+            ResultSet: A SeaResultSet instance for the executed command
+        """
+
+        if session_id.backend_type != BackendType.SEA:
+            raise ValueError("Not a valid SEA session ID")
+
+        sea_session_id = session_id.to_sea_session_id()
+
+        # Convert parameters to StatementParameter objects
+        sea_parameters = []
+        if parameters:
+            for param in parameters:
+                sea_parameters.append(
+                    StatementParameter(
+                        name=param.name,
+                        value=param.value,
+                        type=param.type if hasattr(param, "type") else None,
+                    )
+                )
+
+        format = "ARROW_STREAM" if use_cloud_fetch else "JSON_ARRAY"
+        disposition = "EXTERNAL_LINKS" if use_cloud_fetch else "INLINE"
+        result_compression = "LZ4_FRAME" if lz4_compression else None
+
+        request = ExecuteStatementRequest(
+            warehouse_id=self.warehouse_id,
+            session_id=sea_session_id,
+            statement=operation,
+            disposition=disposition,
+            format=format,
+            wait_timeout="0s" if async_op else "10s",
+            on_wait_timeout="CONTINUE",
+            row_limit=max_rows,
+            parameters=sea_parameters if sea_parameters else None,
+            result_compression=result_compression,
+        )
+
+        response_data = self.http_client._make_request(
+            method="POST", path=self.STATEMENT_PATH, data=request.to_dict()
         )
+        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)
+
+        # Store the command ID in the cursor
+        cursor.active_command_id = command_id
+
+        # If async operation, return and let the client poll for results
+        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,
+                },
+            )
+
+        return self.get_execution_result(command_id, cursor)
 
     def cancel_command(self, command_id: CommandId) -> None:
-        """Not implemented yet."""
-        raise NotImplementedError(
-            "cancel_command is not yet implemented for SEA backend"
+        """
+        Cancel a running command.
+
+        Args:
+            command_id: Command identifier to cancel
+
+        Raises:
+            ValueError: If the command ID is invalid
+        """
+
+        if command_id.backend_type != BackendType.SEA:
+            raise ValueError("Not a valid SEA command ID")
+
+        sea_statement_id = command_id.to_sea_statement_id()
+
+        request = CancelStatementRequest(statement_id=sea_statement_id)
+        self.http_client._make_request(
+            method="POST",
+            path=self.CANCEL_STATEMENT_PATH_WITH_ID.format(sea_statement_id),
+            data=request.to_dict(),
         )
 
     def close_command(self, command_id: CommandId) -> None:
-        """Not implemented yet."""
-        raise NotImplementedError(
-            "close_command is not yet implemented for SEA backend"
+        """
+        Close a command and release resources.
+
+        Args:
+            command_id: Command identifier to close
+
+        Raises:
+            ValueError: If the command ID is invalid
+        """
+
+        if command_id.backend_type != BackendType.SEA:
+            raise ValueError("Not a valid SEA command ID")
+
+        sea_statement_id = command_id.to_sea_statement_id()
+
+        request = CloseStatementRequest(statement_id=sea_statement_id)
+        self.http_client._make_request(
+            method="DELETE",
+            path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
+            data=request.to_dict(),
         )
 
     def get_query_state(self, command_id: CommandId) -> CommandState:
-        """Not implemented yet."""
-        raise NotImplementedError(
-            "get_query_state is not yet implemented for SEA backend"
+        """
+        Get the state of a running query.
+
+        Args:
+            command_id: Command identifier
+
+        Returns:
+            CommandState: The current state of the command
+
+        Raises:
+            ValueError: If the command ID is invalid
+        """
+
+        if command_id.backend_type != BackendType.SEA:
+            raise ValueError("Not a valid SEA command ID")
+
+        sea_statement_id = command_id.to_sea_statement_id()
+
+        request = GetStatementRequest(statement_id=sea_statement_id)
+        response_data = self.http_client._make_request(
+            method="GET",
+            path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
+            data=request.to_dict(),
         )
 
+        # Parse the response
+        response = GetStatementResponse.from_dict(response_data)
+        return response.status.state
+
     def get_execution_result(
         self,
         command_id: CommandId,
         cursor: "Cursor",
-    ):
-        """Not implemented yet."""
-        raise NotImplementedError(
-            "get_execution_result is not yet implemented for SEA backend"
+    ) -> "ResultSet":
+        """
+        Get the result of a command execution.
+
+        Args:
+            command_id: Command identifier
+            cursor: Cursor executing the command
+
+        Returns:
+            ResultSet: A SeaResultSet instance with the execution results
+
+        Raises:
+            ValueError: If the command ID is invalid
+        """
+
+        if command_id.backend_type != BackendType.SEA:
+            raise ValueError("Not a valid SEA command ID")
+
+        sea_statement_id = command_id.to_sea_statement_id()
+
+        # Create the request model
+        request = GetStatementRequest(statement_id=sea_statement_id)
+
+        # Get the statement result
+        response_data = self.http_client._make_request(
+            method="GET",
+            path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
+            data=request.to_dict(),
+        )
+
+        # 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)
+
+        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,
         )
 
     # == Metadata Operations ==