[SPARK-58547][CONNECT] Expose operation IDs for end-to-end request attribution

core/src/main/scala/org/apache/spark/SparkContext.scala

@@ -3157,6 +3157,7 @@ object SparkContext extends Logging {
   private[spark] val SQL_EXECUTION_ID_KEY = "spark.sql.execution.id"
   private[spark] val DATASET_QUERY_EXECUTION_ID_KEY =
     "spark.sql.dataset.queryExecution.id"
+  private[spark] val SPARK_CONNECT_OPERATION_ID_PROPERTY = "spark.connect.operation_id"
 
   /**
    * Executor id for the driver.  In earlier versions of Spark, this was `<driver>`, but this was

python/pyspark/errors/exceptions/connect.py

@@ -53,6 +53,14 @@ class SparkConnectException(PySparkException):
     Exception thrown from Spark Connect.
     """
 
+    @property
+    def operation_id(self) -> Optional[str]:
+        """The Spark Connect ExecutePlan operation ID, when available.
+
+        .. versionadded:: 4.3.0
+        """
+        return getattr(self, "_operation_id", None)
+
 
 def convert_exception(
     info: "ErrorInfo",

python/pyspark/sql/connect/client/core.py

@@ -1239,7 +1239,7 @@ def to_table(
         table, schema, metrics, observed_metrics, _ = self._execute_and_fetch(req, observations)
 
         # Create a query execution object.
-        ei = ExecutionInfo(metrics, observed_metrics)
+        ei = ExecutionInfo(metrics, observed_metrics, req.operation_id)
         assert table is not None
         return table, schema, ei
 
@@ -1275,7 +1275,7 @@ def to_pandas(
             req, observations, selfDestruct == "true"
         )
         assert table is not None
-        ei = ExecutionInfo(metrics, observed_metrics)
+        ei = ExecutionInfo(metrics, observed_metrics, req.operation_id)
 
         schema = schema or from_arrow_schema(table.schema, prefer_timestamp_ntz=True)
         assert schema is not None and isinstance(schema, StructType)
@@ -1418,7 +1418,7 @@ def execute_command(
             req, observations or {}
         )
         # Create a query execution object.
-        ei = ExecutionInfo(metrics, observed_metrics)
+        ei = ExecutionInfo(metrics, observed_metrics, req.operation_id)
         if data is not None:
             return (data.to_pandas(), properties, ei)
         else:
@@ -1535,15 +1535,17 @@ def _execute_plan_request_with_metadata(
                 )
             )
         )
-        if operation_id is not None:
+        if operation_id is None:
+            operation_id = str(uuid.uuid4())
+        else:
             try:
                 uuid.UUID(operation_id, version=4)
             except ValueError as ve:
                 raise PySparkValueError(
                     errorClass="INVALID_OPERATION_UUID_ID",
                     messageParameters={"arg_name": "operation_id", "origin": str(ve)},
                 )
-            req.operation_id = operation_id
+        req.operation_id = operation_id
         self._update_request_with_user_context_extensions(req)
 
         if call_stack_trace := self.__class__._build_call_stack_trace():
@@ -1673,8 +1675,10 @@ def _execute(self, req: pb2.ExecutePlanRequest) -> None:
         """
         logger.debug("Execute")
 
+        operation_id = req.operation_id
         for hook in self._session_hooks:
             req = hook.on_execute_plan(req)
+            req.operation_id = operation_id
 
         def handle_response(b: pb2.ExecutePlanResponse) -> None:
             self._verify_response_integrity(b)
@@ -1703,7 +1707,7 @@ def handle_response(b: pb2.ExecutePlanResponse) -> None:
                             for b in self._stub.ExecutePlan(req, metadata=self._builder.metadata()):
                                 handle_response(b)
         except Exception as error:
-            self._handle_error(error)
+            self._handle_error(error, req.operation_id)
 
     def _execute_and_fetch_as_iterator(
         self,
@@ -1724,8 +1728,10 @@ def _execute_and_fetch_as_iterator(
             # when not at debug log level.
             logger.debug(f"ExecuteAndFetchAsIterator. Request: {self._proto_to_string(req)}")
 
+        operation_id = req.operation_id
         for hook in self._session_hooks:
             req = hook.on_execute_plan(req)
+            req.operation_id = operation_id
 
         num_records = 0
         arrow_batch_chunks_to_assemble: List[bytes] = []
@@ -1932,7 +1938,7 @@ def handle_response(
             self.interrupt_operation(req.operation_id)
             raise kb
         except Exception as error:
-            self._handle_error(error)
+            self._handle_error(error, req.operation_id)
 
     def _execute_and_fetch(
         self,
@@ -2297,7 +2303,7 @@ def clear_user_context_extensions(self) -> None:
         with self.global_user_context_extensions_lock:
             self.global_user_context_extensions = list()
 
-    def _handle_error(self, error: Exception) -> NoReturn:
+    def _handle_error(self, error: Exception, operation_id: Optional[str] = None) -> NoReturn:
         """
         Handle errors that occur during RPC calls.
 
@@ -2318,9 +2324,14 @@ def _handle_error(self, error: Exception) -> NoReturn:
 
         try:
             self.thread_local.inside_error_handling = True
-            if isinstance(error, grpc.RpcError):
-                self._handle_rpc_error(error)
-            raise error
+            try:
+                if isinstance(error, grpc.RpcError):
+                    self._handle_rpc_error(error)
+                raise error
+            except BaseException as handled_error:
+                if operation_id:
+                    handled_error._operation_id = operation_id  # type: ignore[attr-defined]
+                raise
         finally:
             self.thread_local.inside_error_handling = False
 

python/pyspark/sql/metrics.py

@@ -297,10 +297,14 @@ class ExecutionInfo:
     data frame. This value is only set in the data frame if it was executed."""
 
     def __init__(
-        self, metrics: Optional[list[PlanMetrics]], obs: Optional[Sequence[ObservedMetrics]]
+        self,
+        metrics: Optional[list[PlanMetrics]],
+        obs: Optional[Sequence[ObservedMetrics]],
+        operation_id: Optional[str] = None,
     ):
         self._metrics = CollectedMetrics(metrics) if metrics else None
         self._observations = obs if obs else []
+        self._operation_id = operation_id
 
     @property
     def metrics(self) -> Optional[CollectedMetrics]:
@@ -309,3 +313,11 @@ def metrics(self) -> Optional[CollectedMetrics]:
     @property
     def flows(self) -> List[Tuple[str, Dict[str, Any]]]:
         return [(f.name, f.pairs) for f in self._observations]
+
+    @property
+    def operation_id(self) -> Optional[str]:
+        """The Spark Connect ExecutePlan operation ID, when available.
+
+        .. versionadded:: 4.3.0
+        """
+        return self._operation_id

python/pyspark/sql/tests/connect/client/test_client.py

@@ -479,6 +479,33 @@ def on_execute_plan(self, req):
             session.client.close()
             session.stop()
 
+    def test_session_hook_preserves_operation_id(self):
+        class TestHook(RemoteSparkSession.Hook):
+            def __init__(self, _session):
+                pass
+
+            def on_execute_plan(self, req):
+                replacement = proto.ExecutePlanRequest()
+                replacement.CopyFrom(req)
+                replacement.ClearField("operation_id")
+                return replacement
+
+        session = (
+            RemoteSparkSession.builder.remote("sc://foo")._registerHook(TestHook).getOrCreate()
+        )
+        try:
+            mock = MockService(session.client._session_id)
+            session.client._stub = mock
+            session.client.disable_reattachable_execute()
+
+            df = session.range(1)
+            df.collect()
+            self.assertIsNotNone(df.executionInfo)
+            self.assertEqual(mock.req.operation_id, df.executionInfo.operation_id)
+            uuid.UUID(mock.req.operation_id)
+        finally:
+            session.stop()
+
     def test_new_session_preserves_custom_channel_builder(self):
         class CustomChannelBuilder(DefaultChannelBuilder):
             pass
@@ -509,6 +536,14 @@ def test_custom_operation_id(self):
         for resp in client._stub.ExecutePlan(req, metadata=None):
             assert resp.operation_id == "10a4c38e-7e87-40ee-9d6f-60ff0751e63b"
 
+    def test_execute_plan_request_generates_operation_id(self):
+        client = SparkConnectClient("sc://foo/;token=bar", use_reattachable_execute=False)
+        try:
+            req = client._execute_plan_request_with_metadata()
+            uuid.UUID(req.operation_id)
+        finally:
+            client.close()
+
     def test_on_exit_calls_release_and_close_when_enabled(self):
         client = SparkConnectClient("sc://foo/", use_reattachable_execute=False)
         client._release_session_on_exit = True

python/pyspark/sql/tests/connect/test_connect_session.py

@@ -82,6 +82,20 @@ def handler(**kwargs):
         self.spark.sql("select 1").collect()
         self.assertGreaterEqual(len(handler_called), 0)
 
+    @timeout(10)
+    def test_operation_id_in_execution_info_and_exception(self):
+        df = self.spark.sql("select 1")
+        df.collect()
+        self.assertIsNotNone(df.executionInfo)
+        operation_id = df.executionInfo.operation_id
+        self.assertIsNotNone(operation_id)
+        uuid.UUID(operation_id)
+
+        with self.assertRaises(SparkConnectException) as error:
+            self.spark.sql("select raise_error('expected')").collect()
+        self.assertIsNotNone(error.exception.operation_id)
+        uuid.UUID(error.exception.operation_id)
+
     def _check_no_active_session_error(self, e: PySparkException):
         self.check_error(exception=e, errorClass="NO_ACTIVE_SESSION", messageParameters=dict())
 

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala

@@ -78,6 +78,46 @@ class SparkConnectClientSuite extends ConnectFunSuite {
     assert(client.userId == System.getProperty("user.name"))
   }
 
+  test("client generates an operation ID for ExecutePlan requests") {
+    startDummyServer(0)
+    client = SparkConnectClient
+      .builder()
+      .connectionString(s"sc://localhost:${server.getPort}")
+      .disableReattachableExecute()
+      .build()
+
+    val responses = client.execute(buildPlan("select 1")).toSeq
+    val operationId = responses.head.getOperationId
+
+    UUID.fromString(operationId)
+    assert(responses.forall(_.getOperationId == operationId))
+  }
+
+  test("ExecutePlan exceptions expose the client-generated operation ID") {
+    val failingService = new DummySparkConnectService {
+      override def executePlan(
+          request: ExecutePlanRequest,
+          responseObserver: StreamObserver[ExecutePlanResponse]): Unit = {
+        responseObserver.onError(Status.INTERNAL.withDescription("expected").asRuntimeException())
+      }
+    }
+    server = NettyServerBuilder.forPort(0).addService(failingService).build().start()
+    service = failingService
+    client = SparkConnectClient
+      .builder()
+      .connectionString(s"sc://localhost:${server.getPort}")
+      .disableReattachableExecute()
+      .retryPolicy(RetryPolicy(maxRetries = Some(0), canRetry = _ => false, name = "NoRetry"))
+      .build()
+
+    val error = intercept[SparkException] {
+      client.execute(buildPlan("select 1")).foreach(_ => ())
+    }
+    val operationId = SparkConnectClient.getOperationId(error)
+    assert(operationId.isDefined)
+    UUID.fromString(operationId.get)
+  }
+
   test("Placeholder test: Create SparkConnectClient") {
     client = SparkConnectClient.builder().userId("abc123").build()
     assert(client.userId == "abc123")

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala

@@ -52,7 +52,8 @@ private[connect] class CustomSparkConnectBlockingStub(
     grpcExceptionConverter.convert(
       request.getSessionId,
       request.getUserContext,
-      request.getClientType) {
+      request.getClientType,
+      Option(request.getOperationId).filter(_.nonEmpty)) {
       grpcExceptionConverter.convertIterator[ExecutePlanResponse](
         request.getSessionId,
         request.getUserContext,
@@ -62,7 +63,8 @@ private[connect] class CustomSparkConnectBlockingStub(
           r => {
             stubState.responseValidator.wrapIterator(
               CloseableIterator(stub.executePlan(r).asScala))
-          }))
+          }),
+        Option(request.getOperationId).filter(_.nonEmpty))
     }
   }
 
@@ -71,7 +73,8 @@ private[connect] class CustomSparkConnectBlockingStub(
     grpcExceptionConverter.convert(
       request.getSessionId,
       request.getUserContext,
-      request.getClientType) {
+      request.getClientType,
+      Option(request.getOperationId).filter(_.nonEmpty)) {
       grpcExceptionConverter.convertIterator[ExecutePlanResponse](
         request.getSessionId,
         request.getUserContext,
@@ -83,7 +86,8 @@ private[connect] class CustomSparkConnectBlockingStub(
             channel,
             stubState.retryHandler,
             stubState.rpcDeadlines.reattachableExecutePlan,
-            stubState.rpcDeadlines.reattachExecute)))
+            stubState.rpcDeadlines.reattachExecute)),
+        Option(request.getOperationId).filter(_.nonEmpty))
     }
   }
 

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/GrpcExceptionConverter.scala

@@ -64,38 +64,45 @@ private[client] class GrpcExceptionConverter(
       .map(d => grpcStub.withDeadline(Deadline.after(d.toMillis, TimeUnit.MILLISECONDS)))
       .getOrElse(grpcStub)
 
-  def convert[T](sessionId: String, userContext: UserContext, clientType: String)(f: => T): T = {
+  def convert[T](
+      sessionId: String,
+      userContext: UserContext,
+      clientType: String,
+      operationId: Option[String] = None)(f: => T): T = {
     try {
       f
     } catch {
       case e: StatusRuntimeException =>
-        throw toThrowable(e, sessionId, userContext, clientType)
+        val converted = toThrowable(e, sessionId, userContext, clientType)
+        operationId.foreach(SparkConnectClient.attachOperationId(converted, _))
+        throw converted
     }
   }
 
   def convertIterator[T](
       sessionId: String,
       userContext: UserContext,
       clientType: String,
-      iter: CloseableIterator[T]): CloseableIterator[T] = {
+      iter: CloseableIterator[T],
+      operationId: Option[String] = None): CloseableIterator[T] = {
     new WrappedCloseableIterator[T] {
 
       override def innerIterator: Iterator[T] = iter
 
       override def hasNext: Boolean = {
-        convert(sessionId, userContext, clientType) {
+        convert(sessionId, userContext, clientType, operationId) {
           iter.hasNext
         }
       }
 
       override def next(): T = {
-        convert(sessionId, userContext, clientType) {
+        convert(sessionId, userContext, clientType, operationId) {
           iter.next()
         }
       }
 
       override def close(): Unit = {
-        convert(sessionId, userContext, clientType) {
+        convert(sessionId, userContext, clientType, operationId) {
           iter.close()
         }
       }