remove additional call on success

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

Author: varun-edachali-dbx

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

@@ -40,7 +40,6 @@
     DeleteSessionRequest,
     StatementParameter,
     ExecuteStatementResponse,
-    GetStatementResponse,
     CreateSessionResponse,
 )
 
@@ -323,7 +322,7 @@ def _extract_description_from_manifest(
         return columns
 
     def _results_message_to_execute_response(
-        self, response: GetStatementResponse
+        self, response: ExecuteStatementResponse
     ) -> ExecuteResponse:
         """
         Convert a SEA response to an ExecuteResponse and extract result data.
@@ -357,6 +356,28 @@ def _results_message_to_execute_response(
 
         return execute_response
 
+    def _response_to_result_set(
+        self, response: ExecuteStatementResponse, cursor: Cursor
+    ) -> SeaResultSet:
+        """
+        Convert a SEA response to a SeaResultSet.
+        """
+
+        # Create and return a SeaResultSet
+        from databricks.sql.backend.sea.result_set import SeaResultSet
+
+        execute_response = self._results_message_to_execute_response(response)
+
+        return SeaResultSet(
+            connection=cursor.connection,
+            execute_response=execute_response,
+            sea_client=self,
+            result_data=response.result,
+            manifest=response.manifest,
+            buffer_size_bytes=cursor.buffer_size_bytes,
+            arraysize=cursor.arraysize,
+        )
+
     def _check_command_not_in_failed_or_closed_state(
         self, state: CommandState, command_id: CommandId
     ) -> None:
@@ -491,6 +512,10 @@ 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)
+
         self._wait_until_command_done(response)
         return self.get_execution_result(command_id, cursor)
 
@@ -573,7 +598,7 @@ def get_query_state(self, command_id: CommandId) -> CommandState:
         )
 
         # Parse the response
-        response = GetStatementResponse.from_dict(response_data)
+        response = ExecuteStatementResponse.from_dict(response_data)
         return response.status.state
 
     def get_execution_result(
@@ -611,22 +636,9 @@ 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.backend.sea.result_set import SeaResultSet
-
-        execute_response = self._results_message_to_execute_response(response)
+        response = ExecuteStatementResponse.from_dict(response_data)
 
-        return SeaResultSet(
-            connection=cursor.connection,
-            execute_response=execute_response,
-            sea_client=self,
-            result_data=response.result,
-            manifest=response.manifest,
-            buffer_size_bytes=cursor.buffer_size_bytes,
-            arraysize=cursor.arraysize,
-        )
+        return self._response_to_result_set(response, cursor)
 
     # == Metadata Operations ==
 

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

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

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

@@ -124,26 +124,6 @@ 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."""

tests/unit/test_sea_backend.py

@@ -225,7 +225,7 @@ def test_command_execution_sync(
         mock_http_client._make_request.return_value = execute_response
 
         with patch.object(
-            sea_client, "get_execution_result", return_value="mock_result_set"
+            sea_client, "_response_to_result_set", return_value="mock_result_set"
         ) as mock_get_result:
             result = sea_client.execute_command(
                 operation="SELECT 1",
@@ -240,9 +240,6 @@ def test_command_execution_sync(
                 enforce_embedded_schema_correctness=False,
             )
             assert result == "mock_result_set"
-            cmd_id_arg = mock_get_result.call_args[0][0]
-            assert isinstance(cmd_id_arg, CommandId)
-            assert cmd_id_arg.guid == "test-statement-123"
 
         # Test with invalid session ID
         with pytest.raises(ValueError) as excinfo:
@@ -357,7 +354,7 @@ def test_command_execution_advanced(
         mock_http_client._make_request.return_value = execute_response
         param = {"name": "param1", "value": "value1", "type": "STRING"}
 
-        with patch.object(sea_client, "get_execution_result"):
+        with patch.object(sea_client, "_response_to_result_set"):
             sea_client.execute_command(
                 operation="SELECT * FROM table WHERE col = :param1",
                 session_id=sea_session_id,