re-introduce GetStatementResponse

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

Author: varun-edachali-dbx

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

@@ -6,6 +6,7 @@
 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.models.responses import GetStatementResponse
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
     ResultFormat,
@@ -323,7 +324,7 @@ def _extract_description_from_manifest(
         return columns
 
     def _results_message_to_execute_response(
-        self, response: ExecuteStatementResponse
+        self, response: Union[ExecuteStatementResponse, GetStatementResponse]
     ) -> ExecuteResponse:
         """
         Convert a SEA response to an ExecuteResponse and extract result data.
@@ -358,7 +359,9 @@ def _results_message_to_execute_response(
         return execute_response
 
     def _response_to_result_set(
-        self, response: ExecuteStatementResponse, cursor: Cursor
+        self,
+        response: Union[ExecuteStatementResponse, GetStatementResponse],
+        cursor: Cursor,
     ) -> SeaResultSet:
         """
         Convert a SEA response to a SeaResultSet.
@@ -399,22 +402,24 @@ def _check_command_not_in_failed_or_closed_state(
 
     def _wait_until_command_done(
         self, response: ExecuteStatementResponse
-    ) -> ExecuteStatementResponse:
+    ) -> Union[ExecuteStatementResponse, GetStatementResponse]:
         """
         Wait until a command is done.
         """
 
-        state = response.status.state
-        command_id = CommandId.from_sea_statement_id(response.statement_id)
+        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]:
             time.sleep(self.POLL_INTERVAL_SECONDS)
-            response = self._poll_query(command_id)
-            state = response.status.state
+            final_response = self._poll_query(command_id)
+            state = final_response.status.state
 
         self._check_command_not_in_failed_or_closed_state(state, command_id)
 
-        return response
+        return final_response
 
     def execute_command(
         self,
@@ -516,12 +521,11 @@ def execute_command(
         if async_op:
             return None
 
-        if response.status.state == CommandState.SUCCEEDED:
-            # if the response succeeded within the wait_timeout, return the results immediately
-            return self._response_to_result_set(response, cursor)
+        final_response: Union[ExecuteStatementResponse, GetStatementResponse] = response
+        if response.status.state != CommandState.SUCCEEDED:
+            final_response = self._wait_until_command_done(response)
 
-        response = self._wait_until_command_done(response)
-        return self._response_to_result_set(response, cursor)
+        return self._response_to_result_set(final_response, cursor)
 
     def cancel_command(self, command_id: CommandId) -> None:
         """
@@ -573,7 +577,7 @@ def close_command(self, command_id: CommandId) -> None:
             data=request.to_dict(),
         )
 
-    def _poll_query(self, command_id: CommandId) -> ExecuteStatementResponse:
+    def _poll_query(self, command_id: CommandId) -> GetStatementResponse:
         """
         Poll for the current command info.
         """
@@ -591,7 +595,7 @@ def _poll_query(self, command_id: CommandId) -> ExecuteStatementResponse:
             path=self.STATEMENT_PATH_WITH_ID.format(sea_statement_id),
             data=request.to_dict(),
         )
-        response = ExecuteStatementResponse.from_dict(response_data)
+        response = GetStatementResponse.from_dict(response_data)
 
         return response
 

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

@@ -25,6 +25,7 @@
 
 from databricks.sql.backend.sea.models.responses import (
     ExecuteStatementResponse,
+    GetStatementResponse,
     CreateSessionResponse,
 )
 
@@ -46,5 +47,6 @@
     "DeleteSessionRequest",
     # Response models
     "ExecuteStatementResponse",
+    "GetStatementResponse",
     "CreateSessionResponse",
 ]

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

@@ -124,6 +124,26 @@ def from_dict(cls, data: Dict[str, Any]) -> "ExecuteStatementResponse":
         )
 
 
+@dataclass
+class GetStatementResponse:
+    """Representation of the response from getting information about a statement."""
+
+    statement_id: str
+    status: StatementStatus
+    manifest: ResultManifest
+    result: ResultData
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "GetStatementResponse":
+        """Create a GetStatementResponse from a dictionary."""
+        return cls(
+            statement_id=data.get("statement_id", ""),
+            status=_parse_status(data),
+            manifest=_parse_manifest(data),
+            result=_parse_result(data),
+        )
+
+
 @dataclass
 class CreateSessionResponse:
     """Representation of the response from creating a new session."""