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..efbbbd42f9203 --- /dev/null +++ b/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)) + } + } +} 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") + } + } +} 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..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 @@ -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) }