make error more descriptive on command failure

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

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

@@ -5,7 +5,7 @@
 import re
 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.base import ResultManifest, StatementStatus
 from databricks.sql.backend.sea.utils.constants import (
     ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP,
     ResultFormat,
@@ -398,8 +398,9 @@ def _response_to_result_set(
         )
 
     def _check_command_not_in_failed_or_closed_state(
-        self, state: CommandState, command_id: CommandId
+        self, status: StatementStatus, command_id: CommandId
     ) -> None:
+        state = status.state
         if state == CommandState.CLOSED:
             raise DatabaseError(
                 "Command {} unexpectedly closed server side".format(command_id),
@@ -408,8 +409,9 @@ def _check_command_not_in_failed_or_closed_state(
                 },
             )
         if state == CommandState.FAILED:
+            error = status.error
             raise ServerOperationError(
-                "Command {} failed".format(command_id),
+                "Command failed: {} {}".format(error.error_code, error.message),
                 {
                     "operation-id": command_id,
                 },
@@ -423,16 +425,18 @@ def _wait_until_command_done(
         """
 
         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]:
+        while final_response.status.state in [
+            CommandState.PENDING,
+            CommandState.RUNNING,
+        ]:
             time.sleep(self.POLL_INTERVAL_SECONDS)
             final_response = self._poll_query(command_id)
-            state = final_response.status.state
 
-        self._check_command_not_in_failed_or_closed_state(state, command_id)
+        self._check_command_not_in_failed_or_closed_state(
+            final_response.status, command_id
+        )
 
         return final_response
 

tests/unit/test_sea_backend.py

@@ -12,6 +12,7 @@
     SeaDatabricksClient,
     _filter_session_configuration,
 )
+from databricks.sql.backend.sea.models.base import ServiceError, StatementStatus
 from databricks.sql.backend.types import SessionId, CommandId, CommandState, BackendType
 from databricks.sql.parameters.native import IntegerParameter, TDbsqlParameter
 from databricks.sql.thrift_api.TCLIService import ttypes
@@ -408,7 +409,7 @@ def test_command_execution_advanced(
                         async_op=False,
                         enforce_embedded_schema_correctness=False,
                     )
-                assert "Command test-statement-123 failed" in str(excinfo.value)
+                assert "Command failed" in str(excinfo.value)
 
         # Test missing statement ID
         mock_http_client.reset_mock()
@@ -530,18 +531,18 @@ 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
+            StatementStatus(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
+            StatementStatus(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
+                StatementStatus(state=CommandState.CLOSED), sea_command_id
             )
         assert "Command test-statement-123 unexpectedly closed server side" in str(
             excinfo.value
@@ -550,9 +551,13 @@ def test_check_command_state(self, sea_client, sea_command_id):
         # 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
+                StatementStatus(
+                    state=CommandState.FAILED,
+                    error=ServiceError(message="Test error", error_code="TEST_ERROR"),
+                ),
+                sea_command_id,
             )
-        assert "Command test-statement-123 failed" in str(excinfo.value)
+        assert "Command failed" in str(excinfo.value)
 
     def test_utility_methods(self, sea_client):
         """Test utility methods."""