[SPARK-60283][CONNECT] Validate Spark Connect JDBC connections against the server in isValid

sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnection.scala

@@ -57,7 +57,12 @@ class SparkConnectConnection(val url: String, val info: Properties) extends Conn
     }
   }
 
-  override def isValid(timeout: Int): Boolean = !closed && client.isSessionValid
+  override def isValid(timeout: Int): Boolean = {
+    if (timeout < 0) {
+      throw new SQLException("Validation timeout must not be negative.")
+    }
+    !closed && client.validateSession(timeout)
+  }
 
   override def setCatalog(catalog: String): Unit = {
     checkOpen()

sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala

@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connect.client.jdbc
+
+import java.sql.SQLException
+
+import org.apache.spark.sql.connect.client.jdbc.test.JdbcHelper
+import org.apache.spark.sql.connect.test.{ConnectFunSuite, RemoteSparkSession}
+
+class SparkConnectConnectionSuite extends ConnectFunSuite with RemoteSparkSession
+    with JdbcHelper {
+
+  override def jdbcUrl: String = s"jdbc:sc://localhost:$serverPort"
+
+  Seq(0, 5).foreach { timeout =>
+    test(s"isValid detects a server-closed session with timeout $timeout") {
+      withConnection { connection =>
+        val client = connection.asInstanceOf[SparkConnectConnection].client
+
+        // Ensure the server-side session exists before releasing it.
+        connection.getSchema()
+        assert(connection.isValid(timeout))
+
+        client.releaseSession()
+
+        // Local state alone must not reveal that the server closed the session.
+        assert(!connection.isClosed)
+        assert(client.isSessionValid)
+
+        assert(!connection.isValid(timeout))
+      }
+    }
+  }
+
+  test("isValid rejects a negative timeout") {
+    withConnection { connection =>
+      intercept[SQLException] { connection.isValid(-1) }
+    }
+  }
+
+  test("isValid returns false for a closed connection") {
+    withConnection { connection =>
+      connection.close()
+      assert(!connection.isValid(5))
+    }
+  }
+}

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

@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connect.client
+
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicInteger
+
+import io.grpc.Status
+import io.grpc.netty.NettyServerBuilder
+import io.grpc.stub.StreamObserver
+
+import org.apache.spark.connect.proto
+import org.apache.spark.sql.connect.test.ConnectFunSuite
+
+class SparkConnectSessionValidationSuite extends ConnectFunSuite {
+
+  private class ValidationService
+      extends proto.SparkConnectServiceGrpc.SparkConnectServiceImplBase {
+    val calls = new AtomicInteger()
+    @volatile var failure: Option[Status] = None
+    @volatile var respond = true
+    @volatile var lastRequest: proto.ConfigRequest = _
+    @volatile var serverSessionId = UUID.randomUUID().toString
+
+    override def config(
+        request: proto.ConfigRequest,
+        observer: StreamObserver[proto.ConfigResponse]): Unit = {
+      lastRequest = request
+      calls.incrementAndGet()
+      if (respond) {
+        failure match {
+          case Some(status) => observer.onError(status.asRuntimeException())
+          case None =>
+            observer.onNext(
+              proto.ConfigResponse
+                .newBuilder()
+                .setSessionId(request.getSessionId)
+                .setServerSideSessionId(serverSessionId)
+                .build())
+            observer.onCompleted()
+        }
+      }
+    }
+  }
+
+  private def withClient(f: (SparkConnectClient, ValidationService) => Unit): Unit = {
+    val service = new ValidationService
+    val server = NettyServerBuilder.forPort(0).addService(service).build().start()
+    val client = SparkConnectClient.builder().port(server.getPort).build()
+    try {
+      f(client, service)
+    } finally {
+      client.shutdown()
+      server.shutdownNow()
+      assert(server.awaitTermination(5, TimeUnit.SECONDS))
+    }
+  }
+
+  test("validateSession checks the server, including when the timeout is zero") {
+    withClient { (client, service) =>
+      assert(client.validateSession(0))
+      assert(client.validateSession(5))
+      assert(service.calls.get() == 2)
+      assert(service.lastRequest.getSessionId == client.sessionId)
+      assert(service.lastRequest.getClientObservedServerSideSessionId == service.serverSessionId)
+    }
+  }
+
+  test("validateSession returns false when the server closes the session") {
+    withClient { (client, service) =>
+      assert(client.validateSession(5))
+      service.failure = Some(
+        Status.INTERNAL.withDescription(
+          "[INVALID_HANDLE.SESSION_CLOSED] Session was closed. SQLSTATE: 08003"))
+      assert(!client.validateSession(5))
+    }
+  }
+
+  test("validateSession returns false when the server-side session changes") {
+    withClient { (client, service) =>
+      assert(client.validateSession(5))
+      service.serverSessionId = UUID.randomUUID().toString
+      assert(!client.validateSession(5))
+    }
+  }
+
+  test("validateSession bounds validation when the server does not respond") {
+    withClient { (client, service) =>
+      service.respond = false
+      val start = System.nanoTime()
+      assert(!client.validateSession(1))
+      val elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)
+      assert(elapsed >= 500 && elapsed < 5000, s"Validation took $elapsed ms")
+      service.respond = true
+      assert(client.validateSession(5))
+    }
+  }
+
+  test("validateSession returns false promptly when the server is unavailable") {
+    withClient { (client, service) =>
+      service.failure = Some(Status.UNAVAILABLE)
+      val start = System.nanoTime()
+      assert(!client.validateSession(1))
+      val elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)
+      assert(elapsed < 5000, s"Validation took $elapsed ms")
+    }
+  }
+}

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

@@ -129,6 +129,18 @@ private[connect] class CustomSparkConnectBlockingStub(
     }
   }
 
+  def validateSession(request: ConfigRequest, timeout: Int): Unit = {
+    val validationStub = if (timeout > 0) {
+      stub.withDeadlineAfter(timeout.toLong, TimeUnit.SECONDS)
+    } else {
+      stub
+    }
+    // Do not retry or fetch error details: validation must respect its own timeout.
+    stubState.responseValidator.verifyResponse {
+      validationStub.config(request)
+    }
+  }
+
   def interrupt(request: InterruptRequest): InterruptResponse = {
     grpcExceptionConverter.convert(
       request.getSessionId,

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

@@ -185,6 +185,31 @@ private[sql] class SparkConnectClient(
     stubState.responseValidator.isSessionValid
   }
 
+  private[sql] def validateSession(timeout: Int): Boolean = {
+    if (!isSessionValid) {
+      return false
+    }
+    // An empty config read validates the session without running a Spark job.
+    val request = proto.ConfigRequest
+      .newBuilder()
+      .setOperation(
+        proto.ConfigRequest.Operation
+          .newBuilder()
+          .setGet(proto.ConfigRequest.Get.getDefaultInstance))
+      .setSessionId(sessionId)
+      .setClientType(userAgent)
+      .setUserContext(userContext)
+    serverSideSessionId.foreach(request.setClientObservedServerSideSessionId)
+    try {
+      bstub.validateSession(request.build(), timeout)
+      true
+    } catch {
+      case NonFatal(e) =>
+        logDebug("Failed to validate the server-side session.", e)
+        false
+    }
+  }
+
   private[sql] val artifactManager: ArtifactManager = {
     new ArtifactManager(configuration, sessionId, bstub, stub)
   }