adopy _wait_until_command_done

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

Author: varun-edachali-dbx

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

@@ -30,7 +30,7 @@
     BackendType,
     ExecuteResponse,
 )
-from databricks.sql.exc import ServerOperationError
+from databricks.sql.exc import DatabaseError, ServerOperationError
 from databricks.sql.backend.sea.utils.http_client import SeaHttpClient
 from databricks.sql.thrift_api.TCLIService import ttypes
 from databricks.sql.types import SSLOptions
@@ -396,6 +396,42 @@ def _results_message_to_execute_response(self, sea_response, command_id):
 
         return execute_response, result_data_obj, manifest_obj
 
+    def _check_command_not_in_failed_or_closed_state(
+        self, state: CommandState, command_id: CommandId
+    ) -> None:
+        if state == CommandState.CLOSED:
+            raise DatabaseError(
+                "Command {} unexpectedly closed server side".format(command_id),
+                {
+                    "operation-id": command_id,
+                },
+            )
+        if state == CommandState.FAILED:
+            raise ServerOperationError(
+                "Command {} failed".format(command_id),
+                {
+                    "operation-id": command_id,
+                },
+            )
+
+    def _wait_until_command_done(
+        self, response: ExecuteStatementResponse
+    ) -> CommandState:
+        """
+        Wait until a command is done.
+        """
+
+        state = response.status.state
+        command_id = CommandId.from_sea_statement_id(response.statement_id)
+
+        while state in [CommandState.PENDING, CommandState.RUNNING]:
+            time.sleep(self.POLL_INTERVAL_SECONDS)
+            state = self.get_query_state(command_id)
+
+        self._check_command_not_in_failed_or_closed_state(state, command_id)
+
+        return state
+
     def execute_command(
         self,
         operation: str,
@@ -493,24 +529,7 @@ def execute_command(
         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,
-                },
-            )
-
+        self._wait_until_command_done(response)
         return self.get_execution_result(command_id, cursor)
 
     def cancel_command(self, command_id: CommandId) -> None: