From b24270c95c8b013049ebb2576de66d6247553b39 Mon Sep 17 00:00:00 2001 From: Jiwon Park Date: Fri, 16 Oct 2026 23:37:57 +0900 Subject: [PATCH 1/4] [SPARK-60283][CONNECT] Validate Spark Connect JDBC connections against the server in isValid Signed-off-by: Jiwon Park --- .../client/jdbc/SparkConnectConnection.scala | 7 +- .../jdbc/SparkConnectConnectionSuite.scala | 148 ++++++++++++++++++ .../CustomSparkConnectBlockingStub.scala | 12 ++ .../connect/client/SparkConnectClient.scala | 23 +++ 4 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala diff --git a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnection.scala b/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnection.scala index 21b9471bb6069..6142393efec8a 100644 --- a/sql/connect/client/jdbc/src/main/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnection.scala +++ b/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() diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala new file mode 100644 index 0000000000000..02767039d6034 --- /dev/null +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala @@ -0,0 +1,148 @@ +/* + * 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 java.util.{Properties, 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 SparkConnectConnectionSuite 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() + } + } + } + + override def releaseSession( + request: proto.ReleaseSessionRequest, + observer: StreamObserver[proto.ReleaseSessionResponse]): Unit = { + observer.onNext( + proto.ReleaseSessionResponse.newBuilder().setSessionId(request.getSessionId).build()) + observer.onCompleted() + } + } + + private def withConnection(f: (SparkConnectConnection, ValidationService) => Unit): Unit = { + val service = new ValidationService + val server = NettyServerBuilder.forPort(0).addService(service).build().start() + val connection = new SparkConnectConnection( + s"jdbc:sc://localhost:${server.getPort}", new Properties) + try { + f(connection, service) + } finally { + connection.client.shutdown() + server.shutdownNow() + assert(server.awaitTermination(5, TimeUnit.SECONDS)) + } + } + + test("isValid checks the server, including when the timeout is zero") { + withConnection { (connection, service) => + assert(connection.isValid(0)) + assert(connection.isValid(5)) + assert(service.calls.get() == 2) + assert(service.lastRequest.getSessionId == connection.client.sessionId) + assert(service.lastRequest.getClientObservedServerSideSessionId == service.serverSessionId) + } + } + + test("isValid returns false when the server closes the session") { + withConnection { (connection, service) => + assert(connection.isValid(5)) + service.failure = Some(Status.INTERNAL.withDescription( + "[INVALID_HANDLE.SESSION_CLOSED] Session was closed. SQLSTATE: 08003")) + assert(!connection.isValid(5)) + } + } + + test("isValid returns false when the server-side session changes") { + withConnection { (connection, service) => + assert(connection.isValid(5)) + service.serverSessionId = UUID.randomUUID().toString + assert(!connection.isValid(5)) + } + } + + test("isValid bounds validation when the server does not respond") { + withConnection { (connection, service) => + service.respond = false + val start = System.nanoTime() + assert(!connection.isValid(1)) + val elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start) + assert(elapsed >= 500 && elapsed < 5000, s"Validation took $elapsed ms") + service.respond = true + assert(connection.isValid(5)) + } + } + + test("isValid returns false promptly when the server is unavailable") { + withConnection { (connection, service) => + service.failure = Some(Status.UNAVAILABLE) + val start = System.nanoTime() + assert(!connection.isValid(1)) + val elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start) + assert(elapsed < 5000, s"Validation took $elapsed ms") + } + } + + test("isValid rejects a negative timeout") { + withConnection { (connection, service) => + intercept[SQLException] { connection.isValid(-1) } + assert(service.calls.get() == 0) + } + } + + test("isValid returns false without contacting the server for a closed connection") { + withConnection { (connection, service) => + connection.close() + assert(!connection.isValid(5)) + assert(service.calls.get() == 0) + } + } +} diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala index 688029f4883ad..fc04033def7cd 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/CustomSparkConnectBlockingStub.scala +++ b/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, diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala index e60f4a83b6943..1ac2fee683e69 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala @@ -185,6 +185,29 @@ 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(_) => false + } + } + private[sql] val artifactManager: ArtifactManager = { new ArtifactManager(configuration, sessionId, bstub, stub) } From cbc683e4e0bd1e040c15133ae274030c04545a02 Mon Sep 17 00:00:00 2001 From: Jiwon Park Date: Sat, 17 Oct 2026 11:05:12 +0900 Subject: [PATCH 2/4] Separate session validation tests from JDBC integration tests Signed-off-by: Jiwon Park --- .../jdbc/SparkConnectConnectionSuite.scala | 117 ++--------------- .../SparkConnectSessionValidationSuite.scala | 124 ++++++++++++++++++ 2 files changed, 133 insertions(+), 108 deletions(-) create mode 100644 sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectSessionValidationSuite.scala diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala index 02767039d6034..b496c4dad1791 100644 --- a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala @@ -18,131 +18,32 @@ package org.apache.spark.sql.connect.client.jdbc import java.sql.SQLException -import java.util.{Properties, 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.sql.connect.client.jdbc.test.JdbcHelper +import org.apache.spark.sql.connect.test.{ConnectFunSuite, RemoteSparkSession} -import org.apache.spark.connect.proto -import org.apache.spark.sql.connect.test.ConnectFunSuite +class SparkConnectConnectionSuite extends ConnectFunSuite with RemoteSparkSession + with JdbcHelper { -class SparkConnectConnectionSuite 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() - } - } - } - - override def releaseSession( - request: proto.ReleaseSessionRequest, - observer: StreamObserver[proto.ReleaseSessionResponse]): Unit = { - observer.onNext( - proto.ReleaseSessionResponse.newBuilder().setSessionId(request.getSessionId).build()) - observer.onCompleted() - } - } - - private def withConnection(f: (SparkConnectConnection, ValidationService) => Unit): Unit = { - val service = new ValidationService - val server = NettyServerBuilder.forPort(0).addService(service).build().start() - val connection = new SparkConnectConnection( - s"jdbc:sc://localhost:${server.getPort}", new Properties) - try { - f(connection, service) - } finally { - connection.client.shutdown() - server.shutdownNow() - assert(server.awaitTermination(5, TimeUnit.SECONDS)) - } - } + override def jdbcUrl: String = s"jdbc:sc://localhost:$serverPort" test("isValid checks the server, including when the timeout is zero") { - withConnection { (connection, service) => + withConnection { connection => assert(connection.isValid(0)) assert(connection.isValid(5)) - assert(service.calls.get() == 2) - assert(service.lastRequest.getSessionId == connection.client.sessionId) - assert(service.lastRequest.getClientObservedServerSideSessionId == service.serverSessionId) - } - } - - test("isValid returns false when the server closes the session") { - withConnection { (connection, service) => - assert(connection.isValid(5)) - service.failure = Some(Status.INTERNAL.withDescription( - "[INVALID_HANDLE.SESSION_CLOSED] Session was closed. SQLSTATE: 08003")) - assert(!connection.isValid(5)) - } - } - - test("isValid returns false when the server-side session changes") { - withConnection { (connection, service) => - assert(connection.isValid(5)) - service.serverSessionId = UUID.randomUUID().toString - assert(!connection.isValid(5)) - } - } - - test("isValid bounds validation when the server does not respond") { - withConnection { (connection, service) => - service.respond = false - val start = System.nanoTime() - assert(!connection.isValid(1)) - val elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start) - assert(elapsed >= 500 && elapsed < 5000, s"Validation took $elapsed ms") - service.respond = true - assert(connection.isValid(5)) - } - } - - test("isValid returns false promptly when the server is unavailable") { - withConnection { (connection, service) => - service.failure = Some(Status.UNAVAILABLE) - val start = System.nanoTime() - assert(!connection.isValid(1)) - val elapsed = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start) - assert(elapsed < 5000, s"Validation took $elapsed ms") } } test("isValid rejects a negative timeout") { - withConnection { (connection, service) => + withConnection { connection => intercept[SQLException] { connection.isValid(-1) } - assert(service.calls.get() == 0) } } - test("isValid returns false without contacting the server for a closed connection") { - withConnection { (connection, service) => + test("isValid returns false for a closed connection") { + withConnection { connection => connection.close() assert(!connection.isValid(5)) - assert(service.calls.get() == 0) } } } diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectSessionValidationSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectSessionValidationSuite.scala new file mode 100644 index 0000000000000..b258d267f4c09 --- /dev/null +++ b/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") + } + } +} From 9a29d2d781d5984f6fc59d809ca172a0846b7864 Mon Sep 17 00:00:00 2001 From: Jiwon Park Date: Fri, 23 Oct 2026 22:56:03 +0900 Subject: [PATCH 3/4] [SPARK-60283][CONNECT] Test JDBC validation after server session closure Signed-off-by: Jiwon Park --- .../jdbc/SparkConnectConnectionSuite.scala | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala index b496c4dad1791..efbbbd42f9203 100644 --- a/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala +++ b/sql/connect/client/jdbc/src/test/scala/org/apache/spark/sql/connect/client/jdbc/SparkConnectConnectionSuite.scala @@ -27,10 +27,23 @@ class SparkConnectConnectionSuite extends ConnectFunSuite with RemoteSparkSessio override def jdbcUrl: String = s"jdbc:sc://localhost:$serverPort" - test("isValid checks the server, including when the timeout is zero") { - withConnection { connection => - assert(connection.isValid(0)) - assert(connection.isValid(5)) + 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)) + } } } From f0d3e4af39b6cc87f96fada342ce112444660e05 Mon Sep 17 00:00:00 2001 From: Jiwon Park Date: Sat, 24 Oct 2026 16:59:06 +0900 Subject: [PATCH 4/4] [SPARK-60283][CONNECT] Log session validation failures at debug level Signed-off-by: Jiwon Park --- .../apache/spark/sql/connect/client/SparkConnectClient.scala | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala index 1ac2fee683e69..c690c8987a9cd 100644 --- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala +++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala @@ -204,7 +204,9 @@ private[sql] class SparkConnectClient( bstub.validateSession(request.build(), timeout) true } catch { - case NonFatal(_) => false + case NonFatal(e) => + logDebug("Failed to validate the server-side session.", e) + false } }