[SPARK-45798][CONNECT] Assert server-side session ID

### What changes were proposed in this pull request?
Without this patch, when the server would restart because of an abnormal condition, the client would not realize that this be the case. For example, when a driver OOM occurs and the driver is restarted, the client would not realize that the server is restarted and a new session is assigned.

This patch fixes this behavior and asserts that the server side session ID does not change during the connection. If it does change it throws an exception like this:

```
>>> spark.range(10).collect()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/dataframe.py", line 1710, in collect
    table, schema = self._to_table()
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/dataframe.py", line 1722, in _to_table
    table, schema = self._session.client.to_table(query, self._plan.observations)
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 839, in to_table
    table, schema, _, _, _ = self._execute_and_fetch(req, observations)
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 1295, in _execute_and_fetch
    for response in self._execute_and_fetch_as_iterator(req, observations):
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 1273, in _execute_and_fetch_as_iterator
    self._handle_error(error)
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 1521, in _handle_error
    raise error
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 1266, in _execute_and_fetch_as_iterator
    yield from handle_response(b)
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 1193, in handle_response
    self._verify_response_integrity(b)
  File "/Users/martin.grund/Development/spark/python/pyspark/sql/connect/client/core.py", line 1622, in _verify_response_integrity
    raise SparkConnectException(
pyspark.errors.exceptions.connect.SparkConnectException: Received incorrect server side session identifier for request. Please restart Spark Session. (9493c83d-cfa4-488f-9522-838ef3df90bf != c5302e8f-170d-477e-908d-299927b68fd8)
```

### Why are the changes needed?
Stability

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
- Existing tests cover the basic invariant.
- Added new tests.

### Was this patch authored or co-authored using generative AI tooling?
No

Closes #43664 from grundprinzip/SPARK-45798.

Authored-by: Martin Grund <martin.grund@databricks.com>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>

Author: grundprinzip

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

@@ -45,6 +45,7 @@ class ArtifactSuite extends ConnectFunSuite with BeforeAndAfterEach {
   private var retryPolicy: GrpcRetryHandler.RetryPolicy = _
   private var bstub: CustomSparkConnectBlockingStub = _
   private var stub: CustomSparkConnectStub = _
+  private var state: SparkConnectStubState = _
 
   private def startDummyServer(): Unit = {
     service = new DummySparkConnectService()
@@ -58,8 +59,9 @@ class ArtifactSuite extends ConnectFunSuite with BeforeAndAfterEach {
   private def createArtifactManager(): Unit = {
     channel = InProcessChannelBuilder.forName(getClass.getName).directExecutor().build()
     retryPolicy = GrpcRetryHandler.RetryPolicy()
-    bstub = new CustomSparkConnectBlockingStub(channel, retryPolicy)
-    stub = new CustomSparkConnectStub(channel, retryPolicy)
+    state = new SparkConnectStubState(channel, retryPolicy)
+    bstub = new CustomSparkConnectBlockingStub(channel, state)
+    stub = new CustomSparkConnectStub(channel, state)
     artifactManager = new ArtifactManager(Configuration(), "", bstub, stub)
   }
 

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

@@ -197,8 +197,12 @@ message AnalyzePlanRequest {
 
 // Response to performing analysis of the query. Contains relevant metadata to be able to
 // reason about the performance.
+// Next ID: 16
 message AnalyzePlanResponse {
   string session_id = 1;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 15;
 
   oneof result {
     Schema schema = 2;
@@ -317,8 +321,12 @@ message ExecutePlanRequest {
 
 // The response of a query, can be one or more for each request. Responses belonging to the
 // same input query, carry the same `session_id`.
+// Next ID: 16
 message ExecutePlanResponse {
   string session_id = 1;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 15;
 
   // Identifies the ExecutePlan execution.
   // If set by the client in ExecutePlanRequest.operationId, that value is returned.
@@ -492,8 +500,12 @@ message ConfigRequest {
 }
 
 // Response to the config request.
+// Next ID: 5
 message ConfigResponse {
   string session_id = 1;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 4;
 
   // (Optional) The result key-value pairs.
   //
@@ -584,7 +596,17 @@ message AddArtifactsRequest {
 
 // Response to adding an artifact. Contains relevant metadata to verify successful transfer of
 // artifact(s).
+// Next ID: 4
 message AddArtifactsResponse {
+  // Session id in which the AddArtifact was running.
+  string session_id = 2;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 3;
+
+  // The list of artifact(s) seen by the server.
+  repeated ArtifactSummary artifacts = 1;
+
   // Metadata of an artifact.
   message ArtifactSummary {
     string name = 1;
@@ -593,9 +615,6 @@ message AddArtifactsResponse {
     // If false, the client may choose to resend the artifact specified by `name`.
     bool is_crc_successful = 2;
   }
-
-  // The list of artifact(s) seen by the server.
-  repeated ArtifactSummary artifacts = 1;
 }
 
 // Request to get current statuses of artifacts at the server side.
@@ -626,14 +645,20 @@ message ArtifactStatusesRequest {
 }
 
 // Response to checking artifact statuses.
+// Next ID: 4
 message ArtifactStatusesResponse {
+  // Session id in which the ArtifactStatus was running.
+  string session_id = 2;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 3;
+  // A map of artifact names to their statuses.
+  map<string, ArtifactStatus> statuses = 1;
+
   message ArtifactStatus {
     // Exists or not particular artifact at the server.
     bool exists = 1;
   }
-
-  // A map of artifact names to their statuses.
-  map<string, ArtifactStatus> statuses = 1;
 }
 
 message InterruptRequest {
@@ -678,12 +703,17 @@ message InterruptRequest {
   }
 }
 
+// Next ID: 4
 message InterruptResponse {
   // Session id in which the interrupt was running.
   string session_id = 1;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 3;
 
   // Operation ids of the executions which were interrupted.
   repeated string interrupted_ids = 2;
+
 }
 
 message ReattachOptions {
@@ -774,9 +804,13 @@ message ReleaseExecuteRequest {
   }
 }
 
+// Next ID: 4
 message ReleaseExecuteResponse {
   // Session id in which the release was running.
   string session_id = 1;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 3;
 
   // Operation id of the operation on which the release executed.
   // If the operation couldn't be found (because e.g. it was concurrently released), will be unset.
@@ -803,9 +837,13 @@ message ReleaseSessionRequest {
   optional string client_type = 3;
 }
 
+// Next ID: 3
 message ReleaseSessionResponse {
   // Session id of the session on which the release executed.
   string session_id = 1;
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 2;
 }
 
 message FetchErrorDetailsRequest {
@@ -828,8 +866,21 @@ message FetchErrorDetailsRequest {
   optional string client_type = 4;
 }
 
+// Next ID: 5
 message FetchErrorDetailsResponse {
 
+  // Server-side generated idempotency key that the client can use to assert that the server side
+  // session has not changed.
+  string server_side_session_id = 3;
+
+  string session_id = 4;
+
+  // The index of the root error in errors. The field will not be set if the error is not found.
+  optional int32 root_error_idx = 1;
+
+  // A list of errors.
+  repeated Error errors = 2;
+
   message StackTraceElement {
     // The fully qualified name of the class containing the execution point.
     string declaring_class = 1;
@@ -914,12 +965,6 @@ message FetchErrorDetailsResponse {
     // The structured data of a SparkThrowable exception.
     optional SparkThrowable spark_throwable = 5;
   }
-
-  // The index of the root error in errors. The field will not be set if the error is not found.
-  optional int32 root_error_idx = 1;
-
-  // A list of errors.
-  repeated Error errors = 2;
 }
 
 // Main interface for the SparkConnect service.

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

@@ -123,7 +123,12 @@ class ArtifactManager(
       .setSessionId(sessionId)
       .addAllNames(Arrays.asList(artifactName))
       .build()
-    val statuses = bstub.artifactStatus(request).getStatusesMap
+    val response = bstub.artifactStatus(request)
+    if (response.getSessionId != sessionId) {
+      throw new IllegalStateException(
+        s"Session ID mismatch: $sessionId != ${response.getSessionId}")
+    }
+    val statuses = response.getStatusesMap
     if (statuses.containsKey(artifactName)) {
       statuses.get(artifactName).getExists
     } else false
@@ -179,6 +184,9 @@ class ArtifactManager(
     val responseHandler = new StreamObserver[proto.AddArtifactsResponse] {
       private val summaries = mutable.Buffer.empty[ArtifactSummary]
       override def onNext(v: AddArtifactsResponse): Unit = {
+        if (v.getSessionId != sessionId) {
+          throw new IllegalStateException(s"Session ID mismatch: $sessionId != ${v.getSessionId}")
+        }
         v.getArtifactsList.forEach { summary =>
           summaries += summary
         }

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

@@ -24,14 +24,14 @@ import org.apache.spark.connect.proto._
 
 private[connect] class CustomSparkConnectBlockingStub(
     channel: ManagedChannel,
-    retryPolicy: GrpcRetryHandler.RetryPolicy) {
+    stubState: SparkConnectStubState) {
 
   private val stub = SparkConnectServiceGrpc.newBlockingStub(channel)
 
-  private val retryHandler = new GrpcRetryHandler(retryPolicy)
+  private val retryHandler = stubState.retryHandler
 
   // GrpcExceptionConverter with a GRPC stub for fetching error details from server.
-  private val grpcExceptionConverter = new GrpcExceptionConverter(stub)
+  private val grpcExceptionConverter = stubState.exceptionConverter
 
   def executePlan(request: ExecutePlanRequest): CloseableIterator[ExecutePlanResponse] = {
     grpcExceptionConverter.convert(
@@ -44,7 +44,10 @@ private[connect] class CustomSparkConnectBlockingStub(
         request.getClientType,
         retryHandler.RetryIterator[ExecutePlanRequest, ExecutePlanResponse](
           request,
-          r => CloseableIterator(stub.executePlan(r).asScala)))
+          r => {
+            stubState.responseValidator.wrapIterator(
+              CloseableIterator(stub.executePlan(r).asScala))
+          }))
     }
   }
 
@@ -59,7 +62,8 @@ private[connect] class CustomSparkConnectBlockingStub(
         request.getUserContext,
         request.getClientType,
         // Don't use retryHandler - own retry handling is inside.
-        new ExecutePlanResponseReattachableIterator(request, channel, retryPolicy))
+        stubState.responseValidator.wrapIterator(
+          new ExecutePlanResponseReattachableIterator(request, channel, stubState.retryPolicy)))
     }
   }
 
@@ -69,7 +73,9 @@ private[connect] class CustomSparkConnectBlockingStub(
       request.getUserContext,
       request.getClientType) {
       retryHandler.retry {
-        stub.analyzePlan(request)
+        stubState.responseValidator.verifyResponse {
+          stub.analyzePlan(request)
+        }
       }
     }
   }
@@ -80,7 +86,9 @@ private[connect] class CustomSparkConnectBlockingStub(
       request.getUserContext,
       request.getClientType) {
       retryHandler.retry {
-        stub.config(request)
+        stubState.responseValidator.verifyResponse {
+          stub.config(request)
+        }
       }
     }
   }
@@ -91,7 +99,9 @@ private[connect] class CustomSparkConnectBlockingStub(
       request.getUserContext,
       request.getClientType) {
       retryHandler.retry {
-        stub.interrupt(request)
+        stubState.responseValidator.verifyResponse {
+          stub.interrupt(request)
+        }
       }
     }
   }
@@ -102,7 +112,9 @@ private[connect] class CustomSparkConnectBlockingStub(
       request.getUserContext,
       request.getClientType) {
       retryHandler.retry {
-        stub.releaseSession(request)
+        stubState.responseValidator.verifyResponse {
+          stub.releaseSession(request)
+        }
       }
     }
   }
@@ -113,7 +125,9 @@ private[connect] class CustomSparkConnectBlockingStub(
       request.getUserContext,
       request.getClientType) {
       retryHandler.retry {
-        stub.artifactStatus(request)
+        stubState.responseValidator.verifyResponse {
+          stub.artifactStatus(request)
+        }
       }
     }
   }

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

@@ -23,13 +23,13 @@ import org.apache.spark.connect.proto.{AddArtifactsRequest, AddArtifactsResponse
 
 private[client] class CustomSparkConnectStub(
     channel: ManagedChannel,
-    retryPolicy: GrpcRetryHandler.RetryPolicy) {
+    stubState: SparkConnectStubState) {
 
   private val stub = SparkConnectServiceGrpc.newStub(channel)
-  private val retryHandler = new GrpcRetryHandler(retryPolicy)
 
   def addArtifacts(responseObserver: StreamObserver[AddArtifactsResponse])
       : StreamObserver[AddArtifactsRequest] = {
-    retryHandler.RetryStreamObserver(responseObserver, stub.addArtifacts)
+    stubState.responseValidator.wrapStreamObserver(
+      stubState.retryHandler.RetryStreamObserver(responseObserver, stub.addArtifacts))
   }
 }

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

@@ -97,6 +97,10 @@ class ExecutePlanResponseReattachableIterator(
   private[connect] var iter: Option[java.util.Iterator[proto.ExecutePlanResponse]] =
     Some(rawBlockingStub.executePlan(initialRequest))
 
+  // Server side session ID, used to detect if the server side session changed. This is set upon
+  // receiving the first response from the server.
+  private var serverSideSessionId: Option[String] = None
+
   override def innerIterator: Iterator[proto.ExecutePlanResponse] = iter match {
     case Some(it) => it.asScala
     case None =>
@@ -114,10 +118,23 @@ class ExecutePlanResponseReattachableIterator(
 
     try {
       // Get next response, possibly triggering reattach in case of stream error.
-      val ret = retry {
+      val ret: proto.ExecutePlanResponse = retry {
         callIter(_.next())
       }
 
+      // Check if the server-side session state has changed. If this is the case, immediately
+      // abandon execution.
+      serverSideSessionId match {
+        case Some(id) =>
+          if (id != ret.getServerSideSessionId) {
+            throw new IllegalStateException(
+              s"Server side session ID changed. Create a new SparkSession to continue. " +
+                s"(Old: $id, New: ${ret.getServerSideSessionId})")
+          }
+        case None =>
+          serverSideSessionId = Some(ret.getServerSideSessionId)
+      }
+
       // Record last returned response, to know where to restart in case of reattach.
       lastReturnedResponseId = Some(ret.getResponseId)
       if (ret.hasResultComplete) {

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

@@ -22,14 +22,13 @@ import scala.jdk.CollectionConverters._
 import scala.reflect.ClassTag
 
 import com.google.rpc.ErrorInfo
-import io.grpc.StatusRuntimeException
+import io.grpc.{ManagedChannel, StatusRuntimeException}
 import io.grpc.protobuf.StatusProto
 import org.json4s.DefaultFormats
 import org.json4s.jackson.JsonMethods
 
 import org.apache.spark.{QueryContext, QueryContextType, SparkArithmeticException, SparkArrayIndexOutOfBoundsException, SparkDateTimeException, SparkException, SparkIllegalArgumentException, SparkNumberFormatException, SparkRuntimeException, SparkUnsupportedOperationException, SparkUpgradeException}
-import org.apache.spark.connect.proto.{FetchErrorDetailsRequest, FetchErrorDetailsResponse, UserContext}
-import org.apache.spark.connect.proto.SparkConnectServiceGrpc.SparkConnectServiceBlockingStub
+import org.apache.spark.connect.proto.{FetchErrorDetailsRequest, FetchErrorDetailsResponse, SparkConnectServiceGrpc, UserContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.analysis.{NamespaceAlreadyExistsException, NoSuchDatabaseException, NoSuchTableException, TableAlreadyExistsException, TempTableAlreadyExistsException}
@@ -49,10 +48,11 @@ import org.apache.spark.util.ArrayImplicits._
  * the ErrorInfo is missing, the exception will be constructed based on the StatusRuntimeException
  * itself.
  */
-private[client] class GrpcExceptionConverter(grpcStub: SparkConnectServiceBlockingStub)
-    extends Logging {
+private[client] class GrpcExceptionConverter(channel: ManagedChannel) extends Logging {
   import GrpcExceptionConverter._
 
+  val grpcStub = SparkConnectServiceGrpc.newBlockingStub(channel)
+
   def convert[T](sessionId: String, userContext: UserContext, clientType: String)(f: => T): T = {
     try {
       f