remove ExecuteResponse namedtuple, use as backend obj store

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/databricks_client.py

@@ -16,8 +16,6 @@
 
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.backend.types import SessionId, CommandId, CommandState
-from databricks.sql.utils import ExecuteResponse
-from databricks.sql.types import SSLOptions
 
 # Forward reference for type hints
 from typing import TYPE_CHECKING

src/databricks/sql/backend/filters.py

@@ -17,11 +17,8 @@
     TYPE_CHECKING,
 )
 
-# Import SeaResultSet for type checking
-from databricks.sql.backend.sea_result_set import SeaResultSet
-
 if TYPE_CHECKING:
-    from databricks.sql.result_set import ResultSet
+    from databricks.sql.result_set import ResultSet, SeaResultSet
 
 logger = logging.getLogger(__name__)
 
@@ -36,8 +33,8 @@ class ResultSetFilter:
 
     @staticmethod
     def _filter_sea_result_set(
-        result_set: SeaResultSet, filter_func: Callable[[List[Any]], bool]
-    ) -> SeaResultSet:
+        result_set: "SeaResultSet", filter_func: Callable[[List[Any]], bool]
+    ) -> "SeaResultSet":
         """
         Filter a SEA result set using the provided filter function.
 

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

@@ -2,9 +2,11 @@
 import re
 import uuid
 import time
-from typing import Dict, Tuple, List, Optional, Any, Union, TYPE_CHECKING
+from typing import Dict, Set, Tuple, List, Optional, Any, Union, TYPE_CHECKING
 
-from databricks.sql.backend.sea.utils.constants import ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP
+from databricks.sql.backend.sea.utils.constants import (
+    ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
+)
 
 if TYPE_CHECKING:
     from databricks.sql.client import Cursor
@@ -60,6 +62,7 @@ def _filter_session_configuration(
 
     return filtered_session_configuration
 
+
 class SeaDatabricksClient(DatabricksClient):
     """
     Statement Execution API (SEA) implementation of the DatabricksClient interface.
@@ -493,7 +496,7 @@ def get_execution_result(
         )
 
         # Create and return a SeaResultSet
-        from databricks.sql.backend.sea_result_set import SeaResultSet
+        from databricks.sql.result_set import SeaResultSet
 
         return SeaResultSet(
             connection=cursor.connection,

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

@@ -37,8 +37,11 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
                 error_code=error_data.get("error_code"),
             )
 
+        state = CommandState.from_sea_state(status_data.get("state", ""))
+        if state is None:
+            raise ValueError(f"Invalid state: {status_data.get('state', '')}")
         status = StatementStatus(
-            state=CommandState.from_sea_state(status_data.get("state", "")),
+            state=state,
             error=error,
             sql_state=status_data.get("sql_state"),
         )
@@ -72,8 +75,12 @@ def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
                 error_code=error_data.get("error_code"),
             )
 
+        state = CommandState.from_sea_state(status_data.get("state", ""))
+        if state is None:
+            raise ValueError(f"Invalid state: {status_data.get('state', '')}")
+
         status = StatementStatus(
-            state=CommandState.from_sea_state(status_data.get("state", "")),
+            state=state,
             error=error,
             sql_state=status_data.get("sql_state"),
         )

src/databricks/sql/backend/sea_result_set.py

@@ -1,5 +1,6 @@
+from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, Optional, Any
+from typing import Dict, List, Optional, Any, Tuple
 import logging
 
 from databricks.sql.backend.utils import guid_to_hex_id
@@ -80,6 +81,28 @@ def from_thrift_state(
         else:
             return None
 
+    @classmethod
+    def from_sea_state(cls, state: str) -> Optional["CommandState"]:
+        """
+        Map SEA state string to CommandState enum.
+
+        Args:
+            state: SEA state string
+
+        Returns:
+            CommandState: The corresponding CommandState enum value
+        """
+        state_mapping = {
+            "PENDING": cls.PENDING,
+            "RUNNING": cls.RUNNING,
+            "SUCCEEDED": cls.SUCCEEDED,
+            "FAILED": cls.FAILED,
+            "CLOSED": cls.CLOSED,
+            "CANCELED": cls.CANCELLED,
+        }
+
+        return state_mapping.get(state, None)
+
 
 class BackendType(Enum):
     """
@@ -371,3 +394,19 @@ def to_hex_guid(self) -> str:
             return guid_to_hex_id(self.guid)
         else:
             return str(self.guid)
+
+
+@dataclass
+class ExecuteResponse:
+    """Response from executing a SQL command."""
+
+    command_id: CommandId
+    status: CommandState
+    description: Optional[
+        List[Tuple[str, str, None, None, Optional[int], Optional[int], bool]]
+    ] = None
+    has_more_rows: bool = False
+    results_queue: Optional[Any] = None
+    has_been_closed_server_side: bool = False
+    lz4_compressed: bool = True
+    is_staging_operation: bool = False

src/databricks/sql/backend/types.py

@@ -24,7 +24,6 @@
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.utils import (
-    ExecuteResponse,
     ParamEscaper,
     inject_parameters,
     transform_paramstyle,