[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python

connector/connect/common/src/main/protobuf/spark/connect/commands.proto

@@ -259,15 +259,21 @@ message StreamingQueryCommand {
     bool process_all_available = 6;
     // explain() API. Returns logical and physical plans.
     ExplainCommand explain = 7;
-
-    // TODO(SPARK-42960) Add more commands: await_termination(), exception() etc.
+    // exception() API. Returns the exception in the query if any.
+    bool exception = 8;
+    // awaitTermination() API. Waits for the termination of the query.
+    AwaitTerminationCommand await_termination = 9;
   }
 
   message ExplainCommand {
     // TODO: Consider reusing Explain from AnalyzePlanRequest message.
     //       We can not do this right now since it base.proto imports this file.
     bool extended = 1;
   }
+
+  message AwaitTerminationCommand {
+    int64 timeout_ms = 1;
+  }
 }
 
 // Response for commands on a streaming query.
@@ -279,6 +285,8 @@ message StreamingQueryCommandResult {
     StatusResult status = 2;
     RecentProgressResult recent_progress = 3;
     ExplainResult explain = 4;
+    ExceptionResult exception = 5;
+    AwaitTerminationResult await_termination = 6;
   }
 
   message StatusResult {
@@ -298,6 +306,16 @@ message StreamingQueryCommandResult {
     // Logical and physical plans as string
     string result = 1;
   }
+
+  message ExceptionResult {
+    // Exception as string
+    bool has_exception = 1;
+    optional string error_message = 2;
+  }
+
+  message AwaitTerminationResult {
+    bool terminated = 1;
+  }
 }
 
 // Command to get the output of 'SparkContext.resources'

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -52,7 +52,7 @@ import org.apache.spark.sql.connect.artifact.SparkConnectArtifactManager
 import org.apache.spark.sql.connect.common.{DataTypeProtoConverter, InvalidPlanInput, LiteralValueProtoConverter, StorageLevelProtoConverter, UdfPacket}
 import org.apache.spark.sql.connect.config.Connect.CONNECT_GRPC_ARROW_MAX_BATCH_SIZE
 import org.apache.spark.sql.connect.plugin.SparkConnectPluginRegistry
-import org.apache.spark.sql.connect.service.SparkConnectStreamHandler
+import org.apache.spark.sql.connect.service.{SparkConnectService, SparkConnectStreamHandler}
 import org.apache.spark.sql.errors.QueryCompilationErrors
 import org.apache.spark.sql.execution.QueryExecution
 import org.apache.spark.sql.execution.arrow.ArrowConverters
@@ -2221,6 +2221,31 @@ class SparkConnectPlanner(val session: SparkSession) {
           .build()
         respBuilder.setExplain(explain)
 
+      case StreamingQueryCommand.CommandCase.EXCEPTION =>
+        val result = query.exception
+        val exception = result match {
+          case Some(e) =>
+            StreamingQueryCommandResult.ExceptionResult
+              .newBuilder()
+              .setHasException(true)
+              .setErrorMessage(SparkConnectService.extractErrorMessage(e))
+              .build()
+          case None =>
+            StreamingQueryCommandResult.ExceptionResult
+              .newBuilder()
+              .setHasException(false)
+              .build()
+        }
+        respBuilder.setException(exception)
+
+      case StreamingQueryCommand.CommandCase.AWAIT_TERMINATION =>
+        val terminated = query.awaitTermination(command.getAwaitTermination.getTimeoutMs)
+        val terminatedResult = StreamingQueryCommandResult.AwaitTerminationResult
+          .newBuilder()
+          .setTerminated(terminated)
+          .build()
+        respBuilder.setAwaitTermination(terminatedResult)
+
       case StreamingQueryCommand.CommandCase.COMMAND_NOT_SET =>
         throw new IllegalArgumentException("Missing command in StreamingQueryCommand")
     }

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

@@ -56,6 +56,8 @@ class SparkConnectService(debug: Boolean)
     extends proto.SparkConnectServiceGrpc.SparkConnectServiceImplBase
     with Logging {
 
+  import SparkConnectService._
+
   private def allClasses(cl: Class[_]): Seq[Class[_]] = {
     val classes = ArrayBuffer.empty[Class[_]]
     if (cl != null && !cl.equals(classOf[java.lang.Object])) {
@@ -85,7 +87,7 @@ class SparkConnectService(debug: Boolean)
             .setDomain("org.apache.spark")
             .putMetadata("classes", compact(render(allClasses(st.getClass).map(_.getName))))
             .build()))
-      .setMessage(StringUtils.abbreviate(st.getMessage, 2048))
+      .setMessage(extractErrorMessage(st))
       .build()
   }
 
@@ -294,4 +296,8 @@ object SparkConnectService {
       }
     }
   }
+
+  def extractErrorMessage(st: Throwable): String = {
+    StringUtils.abbreviate(st.getMessage, 2048)
+  }
 }

python/pyspark/sql/connect/proto/commands_pb2.pyi

@@ -876,13 +876,29 @@ class StreamingQueryCommand(google.protobuf.message.Message):
             self, field_name: typing_extensions.Literal["extended", b"extended"]
         ) -> None: ...
 
+    class AwaitTerminationCommand(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        TIMEOUT_MS_FIELD_NUMBER: builtins.int
+        timeout_ms: builtins.int
+        def __init__(
+            self,
+            *,
+            timeout_ms: builtins.int = ...,
+        ) -> None: ...
+        def ClearField(
+            self, field_name: typing_extensions.Literal["timeout_ms", b"timeout_ms"]
+        ) -> None: ...
+
     QUERY_ID_FIELD_NUMBER: builtins.int
     STATUS_FIELD_NUMBER: builtins.int
     LAST_PROGRESS_FIELD_NUMBER: builtins.int
     RECENT_PROGRESS_FIELD_NUMBER: builtins.int
     STOP_FIELD_NUMBER: builtins.int
     PROCESS_ALL_AVAILABLE_FIELD_NUMBER: builtins.int
     EXPLAIN_FIELD_NUMBER: builtins.int
+    EXCEPTION_FIELD_NUMBER: builtins.int
+    AWAIT_TERMINATION_FIELD_NUMBER: builtins.int
     @property
     def query_id(self) -> global___StreamingQueryInstanceId:
         """(Required) Query instance. See `StreamingQueryInstanceId`."""
@@ -899,6 +915,11 @@ class StreamingQueryCommand(google.protobuf.message.Message):
     @property
     def explain(self) -> global___StreamingQueryCommand.ExplainCommand:
         """explain() API. Returns logical and physical plans."""
+    exception: builtins.bool
+    """exception() API. Returns the exception in the query if any."""
+    @property
+    def await_termination(self) -> global___StreamingQueryCommand.AwaitTerminationCommand:
+        """awaitTermination() API. Waits for the termination of the query."""
     def __init__(
         self,
         *,
@@ -909,12 +930,18 @@ class StreamingQueryCommand(google.protobuf.message.Message):
         stop: builtins.bool = ...,
         process_all_available: builtins.bool = ...,
         explain: global___StreamingQueryCommand.ExplainCommand | None = ...,
+        exception: builtins.bool = ...,
+        await_termination: global___StreamingQueryCommand.AwaitTerminationCommand | None = ...,
     ) -> None: ...
     def HasField(
         self,
         field_name: typing_extensions.Literal[
+            "await_termination",
+            b"await_termination",
             "command",
             b"command",
+            "exception",
+            b"exception",
             "explain",
             b"explain",
             "last_progress",
@@ -934,8 +961,12 @@ class StreamingQueryCommand(google.protobuf.message.Message):
     def ClearField(
         self,
         field_name: typing_extensions.Literal[
+            "await_termination",
+            b"await_termination",
             "command",
             b"command",
+            "exception",
+            b"exception",
             "explain",
             b"explain",
             "last_progress",
@@ -955,7 +986,14 @@ class StreamingQueryCommand(google.protobuf.message.Message):
     def WhichOneof(
         self, oneof_group: typing_extensions.Literal["command", b"command"]
     ) -> typing_extensions.Literal[
-        "status", "last_progress", "recent_progress", "stop", "process_all_available", "explain"
+        "status",
+        "last_progress",
+        "recent_progress",
+        "stop",
+        "process_all_available",
+        "explain",
+        "exception",
+        "await_termination",
     ] | None: ...
 
 global___StreamingQueryCommand = StreamingQueryCommand
@@ -1033,10 +1071,61 @@ class StreamingQueryCommandResult(google.protobuf.message.Message):
             self, field_name: typing_extensions.Literal["result", b"result"]
         ) -> None: ...
 
+    class ExceptionResult(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        HAS_EXCEPTION_FIELD_NUMBER: builtins.int
+        ERROR_MESSAGE_FIELD_NUMBER: builtins.int
+        has_exception: builtins.bool
+        """Exception as string"""
+        error_message: builtins.str
+        def __init__(
+            self,
+            *,
+            has_exception: builtins.bool = ...,
+            error_message: builtins.str | None = ...,
+        ) -> None: ...
+        def HasField(
+            self,
+            field_name: typing_extensions.Literal[
+                "_error_message", b"_error_message", "error_message", b"error_message"
+            ],
+        ) -> builtins.bool: ...
+        def ClearField(
+            self,
+            field_name: typing_extensions.Literal[
+                "_error_message",
+                b"_error_message",
+                "error_message",
+                b"error_message",
+                "has_exception",
+                b"has_exception",
+            ],
+        ) -> None: ...
+        def WhichOneof(
+            self, oneof_group: typing_extensions.Literal["_error_message", b"_error_message"]
+        ) -> typing_extensions.Literal["error_message"] | None: ...
+
+    class AwaitTerminationResult(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+        TERMINATED_FIELD_NUMBER: builtins.int
+        terminated: builtins.bool
+        def __init__(
+            self,
+            *,
+            terminated: builtins.bool = ...,
+        ) -> None: ...
+        def ClearField(
+            self, field_name: typing_extensions.Literal["terminated", b"terminated"]
+        ) -> None: ...
+
     QUERY_ID_FIELD_NUMBER: builtins.int
     STATUS_FIELD_NUMBER: builtins.int
     RECENT_PROGRESS_FIELD_NUMBER: builtins.int
     EXPLAIN_FIELD_NUMBER: builtins.int
+    EXCEPTION_FIELD_NUMBER: builtins.int
+    AWAIT_TERMINATION_FIELD_NUMBER: builtins.int
     @property
     def query_id(self) -> global___StreamingQueryInstanceId:
         """(Required) Query instance id. See `StreamingQueryInstanceId`."""
@@ -1046,17 +1135,27 @@ class StreamingQueryCommandResult(google.protobuf.message.Message):
     def recent_progress(self) -> global___StreamingQueryCommandResult.RecentProgressResult: ...
     @property
     def explain(self) -> global___StreamingQueryCommandResult.ExplainResult: ...
+    @property
+    def exception(self) -> global___StreamingQueryCommandResult.ExceptionResult: ...
+    @property
+    def await_termination(self) -> global___StreamingQueryCommandResult.AwaitTerminationResult: ...
     def __init__(
         self,
         *,
         query_id: global___StreamingQueryInstanceId | None = ...,
         status: global___StreamingQueryCommandResult.StatusResult | None = ...,
         recent_progress: global___StreamingQueryCommandResult.RecentProgressResult | None = ...,
         explain: global___StreamingQueryCommandResult.ExplainResult | None = ...,
+        exception: global___StreamingQueryCommandResult.ExceptionResult | None = ...,
+        await_termination: global___StreamingQueryCommandResult.AwaitTerminationResult | None = ...,
     ) -> None: ...
     def HasField(
         self,
         field_name: typing_extensions.Literal[
+            "await_termination",
+            b"await_termination",
+            "exception",
+            b"exception",
             "explain",
             b"explain",
             "query_id",
@@ -1072,6 +1171,10 @@ class StreamingQueryCommandResult(google.protobuf.message.Message):
     def ClearField(
         self,
         field_name: typing_extensions.Literal[
+            "await_termination",
+            b"await_termination",
+            "exception",
+            b"exception",
             "explain",
             b"explain",
             "query_id",
@@ -1086,7 +1189,9 @@ class StreamingQueryCommandResult(google.protobuf.message.Message):
     ) -> None: ...
     def WhichOneof(
         self, oneof_group: typing_extensions.Literal["result_type", b"result_type"]
-    ) -> typing_extensions.Literal["status", "recent_progress", "explain"] | None: ...
+    ) -> typing_extensions.Literal[
+        "status", "recent_progress", "explain", "exception", "await_termination"
+    ] | None: ...
 
 global___StreamingQueryCommandResult = StreamingQueryCommandResult
 

python/pyspark/sql/connect/streaming/query.py

@@ -16,6 +16,7 @@
 #
 
 import json
+import time
 import sys
 from typing import TYPE_CHECKING, Any, cast, Dict, List, Optional
 
@@ -24,6 +25,9 @@
 from pyspark.sql.streaming.query import (
     StreamingQuery as PySparkStreamingQuery,
 )
+from pyspark.errors.exceptions.connect import (
+    StreamingQueryException as CapturedStreamingQueryException,
+)
 
 __all__ = [
     "StreamingQuery",  # TODO(SPARK-43032): "StreamingQueryManager"
@@ -66,11 +70,35 @@ def isActive(self) -> bool:
 
     isActive.__doc__ = PySparkStreamingQuery.isActive.__doc__
 
-    # TODO (SPARK-42960): Implement and uncomment the doc
+    def _execute_await_termination_cmd(self, timeout: int = 10) -> bool:
+        cmd = pb2.StreamingQueryCommand()
+        cmd.await_termination.timeout_ms = timeout
+        terminated = self._execute_streaming_query_cmd(cmd).await_termination.terminated
+        return terminated
+
+    def _await_termination(self, timeoutMs: Optional[int]) -> Optional[bool]:
+        terminated = False
+        if timeoutMs is None:
+            while not terminated:
+                # When no timeout is set, query the server every 10ms until query terminates
+                terminated = self._execute_await_termination_cmd()
+        else:
+            reqTimeoutMs = min(timeoutMs, 10)
+            while timeoutMs > 0 and not terminated:
+                # When timeout is set, query the server every reqTimeoutMs until query terminates or timout
+                start = time.time()
+                terminated = self._execute_await_termination_cmd(reqTimeoutMs)
+                end = time.time()
+                timeoutMs -= (end - start) * 1000
+            return terminated
+
     def awaitTermination(self, timeout: Optional[int] = None) -> Optional[bool]:
-        raise NotImplementedError()
+        if timeout is not None:
+            if not isinstance(timeout, (int, float)) or timeout < 0:
+                raise ValueError("timeout must be a positive integer or float. Got %s" % timeout)
+        return self._await_termination(int(timeout * 1000))
 
-    # awaitTermination.__doc__ = PySparkStreamingQuery.awaitTermination.__doc__
+    awaitTermination.__doc__ = PySparkStreamingQuery.awaitTermination.__doc__
 
     @property
     def status(self) -> Dict[str, Any]:
@@ -127,9 +155,14 @@ def explain(self, extended: bool = False) -> None:
 
     explain.__doc__ = PySparkStreamingQuery.explain.__doc__
 
-    # TODO (SPARK-42960): Implement and uncomment the doc
     def exception(self) -> Optional[StreamingQueryException]:
-        raise NotImplementedError()
+        cmd = pb2.StreamingQueryCommand()
+        cmd.exception = True
+        exception = self._execute_streaming_query_cmd(cmd).exception
+        if not exception.has_exception:
+            return None
+        else:
+            return CapturedStreamingQueryException(exception.error_message)
 
     exception.__doc__ = PySparkStreamingQuery.exception.__doc__