diff --git a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionImpl.scala b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionImpl.scala index e3bf76e7298..a5d227588a9 100644 --- a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionImpl.scala +++ b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionImpl.scala @@ -16,9 +16,11 @@ */ package org.apache.kyuubi.engine.dataagent.session +import scala.util.control.NonFatal + import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiSQLException} import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY -import org.apache.kyuubi.session.{AbstractSession, SessionHandle, SessionManager} +import org.apache.kyuubi.session.{AbstractSession, SessionHandle} import org.apache.kyuubi.shaded.hive.service.rpc.thrift.{TGetInfoType, TGetInfoValue, TProtocolVersion} class DataAgentSessionImpl( @@ -27,20 +29,20 @@ class DataAgentSessionImpl( password: String, ipAddress: String, conf: Map[String, String], - sessionManager: SessionManager) - extends AbstractSession(protocol, user, password, ipAddress, conf, sessionManager) { + dataAgentSessionManager: DataAgentSessionManager) + extends AbstractSession(protocol, user, password, ipAddress, conf, dataAgentSessionManager) { override val handle: SessionHandle = conf.get(KYUUBI_SESSION_HANDLE_KEY).map(SessionHandle.fromUUID).getOrElse(SessionHandle()) - private val dataAgentProvider = - sessionManager.asInstanceOf[DataAgentSessionManager].dataAgentProvider + private val dataAgentProvider = dataAgentSessionManager.dataAgentProvider override def open(): Unit = { info(s"Starting to open data agent session.") dataAgentProvider.open(handle.identifier.toString, user) try { super.open() + dataAgentSessionManager.registerRoute(handle.identifier.toString, user) } catch { case e: Throwable => try { @@ -72,6 +74,11 @@ class DataAgentSessionImpl( } override def close(): Unit = { + try { + dataAgentSessionManager.unregisterRoute(handle.identifier.toString) + } catch { + case NonFatal(e) => warn("Failed to unregister data agent session route", e) + } try { dataAgentProvider.close(handle.identifier.toString) } finally { diff --git a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionManager.scala b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionManager.scala index 30b6d5ba6ba..574f9b37fff 100644 --- a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionManager.scala +++ b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/session/DataAgentSessionManager.scala @@ -18,11 +18,13 @@ package org.apache.kyuubi.engine.dataagent.session import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf.ENGINE_SHARE_LEVEL -import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SESSION_HANDLE_KEY +import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_SERVER_HA_NAMESPACE_KEY import org.apache.kyuubi.engine.ShareLevel -import org.apache.kyuubi.engine.dataagent.DataAgentEngine +import org.apache.kyuubi.engine.dataagent.{DataAgentEngine, DataAgentTBinaryFrontendService} import org.apache.kyuubi.engine.dataagent.operation.DataAgentOperationManager import org.apache.kyuubi.engine.dataagent.provider.DataAgentProvider +import org.apache.kyuubi.ha.HighAvailabilityConf.{HA_ENGINE_REF_ID, HA_NAMESPACE} +import org.apache.kyuubi.ha.client.{DataAgentSessionRoute, EngineServiceDiscovery} import org.apache.kyuubi.operation.OperationManager import org.apache.kyuubi.session.{Session, SessionHandle, SessionManager} import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TProtocolVersion @@ -50,10 +52,34 @@ class DataAgentSessionManager(name: String) password: String, ipAddress: String, conf: Map[String, String]): Session = { - conf.get(KYUUBI_SESSION_HANDLE_KEY).map(SessionHandle.fromUUID) - .flatMap(getSessionOption).getOrElse { - new DataAgentSessionImpl(protocol, user, password, ipAddress, conf, this) - } + new DataAgentSessionImpl(protocol, user, password, ipAddress, conf, this) + } + + private def discoveryClient = DataAgentEngine.currentEngine + .flatMap(_.frontendServices.collectFirst { + case frontend: DataAgentTBinaryFrontendService => frontend + }).flatMap(_.discoveryService.collect { + case service: EngineServiceDiscovery => service.discoveryClient + }) + + private[dataagent] def registerRoute(sessionId: String, user: String): Unit = { + for { + serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY) + engineRefId <- conf.get(HA_ENGINE_REF_ID) + discovery <- discoveryClient + } { + val route = DataAgentSessionRoute(conf.get(HA_NAMESPACE), engineRefId, user) + DataAgentSessionRoute.register(discovery, serverSpace, sessionId, route) + } + } + + private[dataagent] def unregisterRoute(sessionId: String): Unit = { + for { + serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY) + discovery <- discoveryClient + } { + DataAgentSessionRoute.unregister(discovery, serverSpace, sessionId) + } } override def stop(): Unit = { diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala index fb1385a3ac2..eb4d009d72c 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiReservedKeys.scala @@ -21,6 +21,7 @@ object KyuubiReservedKeys { final val KYUUBI_CLIENT_IP_KEY = "kyuubi.client.ipAddress" final val KYUUBI_CLIENT_VERSION_KEY = "kyuubi.client.version" final val KYUUBI_SERVER_IP_KEY = "kyuubi.server.ipAddress" + final val KYUUBI_SERVER_HA_NAMESPACE_KEY = "kyuubi.server.ha.namespace" final val KYUUBI_SESSION_USER_KEY = "kyuubi.session.user" final val KYUUBI_SESSION_SIGN_PUBLICKEY = "kyuubi.session.sign.publickey" final val KYUUBI_SESSION_USER_SIGN = "kyuubi.session.user.sign" diff --git a/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DataAgentSessionRoute.scala b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DataAgentSessionRoute.scala new file mode 100644 index 00000000000..c07a5d6e7d8 --- /dev/null +++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DataAgentSessionRoute.scala @@ -0,0 +1,103 @@ +/* + * 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.kyuubi.ha.client + +import java.nio.charset.StandardCharsets +import java.util.Base64 + +import scala.util.control.NonFatal + +case class DataAgentSessionRoute(engineSpace: String, engineRefId: String, user: String) + +object DataAgentSessionRoute { + private val ROOT_SUFFIX = "DATA_AGENT_sessions" + + private def root(serverSpace: String): String = s"${serverSpace}_$ROOT_SUFFIX" + + private def path(serverSpace: String, sessionId: String): String = + DiscoveryPaths.makePath(root(serverSpace), sessionId) + + def register( + discovery: DiscoveryClient, + serverSpace: String, + sessionId: String, + route: DataAgentSessionRoute): Unit = { + val routePath = path(serverSpace, sessionId) + val entry = encode(route) + children(discovery, routePath) match { + case Nil => + val entryPath = DiscoveryPaths.makePath(routePath, entry) + try discovery.create(entryPath, "PERSISTENT") + catch { + case NonFatal(_) if discovery.pathExists(entryPath) => + } + case existing :: Nil if existing == entry => + case _ => + throw new IllegalStateException(s"A different route exists for session $sessionId") + } + } + + def find( + discovery: DiscoveryClient, + serverSpace: String, + sessionId: String): Option[DataAgentSessionRoute] = { + val routePath = path(serverSpace, sessionId) + children(discovery, routePath) match { + case Nil => None + case entry :: Nil => Some(decode(entry)) + case _ => throw new IllegalStateException(s"Multiple routes found for session $sessionId") + } + } + + def unregister( + discovery: DiscoveryClient, + serverSpace: String, + sessionId: String): Unit = { + val routePath = path(serverSpace, sessionId) + try { + discovery.delete(routePath, deleteChildren = true) + } catch { + case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) => + } + } + + private def encode(route: DataAgentSessionRoute): String = { + Seq(route.engineSpace, route.engineRefId, route.user) + .map(encodeValue) + .mkString(".") + } + + private def encodeValue(value: String): String = + Base64.getUrlEncoder.withoutPadding() + .encodeToString(value.getBytes(StandardCharsets.UTF_8)) + + private def children(discovery: DiscoveryClient, routePath: String): List[String] = { + try discovery.getChildren(routePath) + catch { + case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) => Nil + } + } + + private def decode(entry: String): DataAgentSessionRoute = { + val values = entry.split("\\.", -1) + require(values.length == 3, "Invalid Data Agent session route") + def decodeValue(value: String): String = + new String(Base64.getUrlDecoder.decode(value), StandardCharsets.UTF_8) + DataAgentSessionRoute(decodeValue(values(0)), decodeValue(values(1)), decodeValue(values(2))) + } +} diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DiscoveryClientTests.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DiscoveryClientTests.scala index 53c0586f5a6..426fbbe4555 100644 --- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DiscoveryClientTests.scala +++ b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DiscoveryClientTests.scala @@ -170,4 +170,38 @@ trait DiscoveryClientTests extends KyuubiFunSuite { assert(data == dataFromGet) } } + + test("data agent session route lifecycle") { + testDataAgentSessionRouteLifecycle() + } + + protected def testDataAgentSessionRouteLifecycle(): Unit = { + withDiscoveryClient(conf) { discoveryClient => + val serverSpace = "server-space" + val sessionId = "session-id" + val route = DataAgentSessionRoute( + "engine space/with.delimiters", + "engine-ref:id/with.delimiters", + "alice+route@example.com") + + assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty) + DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route) + DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route) + assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).contains(route)) + + val conflictingRoute = route.copy(engineRefId = "another-engine-ref-id") + intercept[IllegalStateException] { + DataAgentSessionRoute.register( + discoveryClient, + serverSpace, + sessionId, + conflictingRoute) + } + assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).contains(route)) + + DataAgentSessionRoute.unregister(discoveryClient, serverSpace, sessionId) + DataAgentSessionRoute.unregister(discoveryClient, serverSpace, sessionId) + assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty) + } + } } diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClientSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClientSuite.scala index 70e2a86fbef..56d4c808745 100644 --- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClientSuite.scala +++ b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/etcd/EtcdDiscoveryClientSuite.scala @@ -110,4 +110,8 @@ class EtcdDiscoveryClientSuite extends DiscoveryClientTests { assert(!discoveryClient.pathExists(path)) } } + + test("etcd data agent session route lifecycle") { + testDataAgentSessionRouteLifecycle() + } } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/client/KyuubiSyncThriftClient.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/client/KyuubiSyncThriftClient.scala index c36e9ec06c7..fef2adee24e 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/client/KyuubiSyncThriftClient.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/client/KyuubiSyncThriftClient.scala @@ -62,6 +62,18 @@ class KyuubiSyncThriftClient private ( // Visible for testing. private[kyuubi] def remoteSessionHandle: TSessionHandle = _remoteSessionHandle + private[kyuubi] def useExistingSessionHandle(sessionHandle: SessionHandle): Unit = { + require(_remoteSessionHandle == null, "A session handle is already set") + _remoteSessionHandle = sessionHandle.toTSessionHandle + } + + private[kyuubi] def closeTransport(): Unit = { + Seq(protocol).union(engineAliveProbeProtocol.toSeq).foreach { tProtocol => + if (tProtocol.getTransport.isOpen) tProtocol.getTransport.close() + } + shutdownAsyncRequestExecutor() + } + @volatile private var _aliveProbeSessionHandle: TSessionHandle = _ @volatile private var _remoteEngineBroken: Boolean = false private[kyuubi] def remoteEngineBroken: Boolean = _remoteEngineBroken @@ -514,4 +526,16 @@ private[kyuubi] object KyuubiSyncThriftClient extends Logging { aliveProbeInterval, aliveTimeout) } + + def createClientWithExistingSessionHandle( + user: String, + password: String, + host: String, + port: Int, + conf: KyuubiConf, + sessionHandle: SessionHandle): KyuubiSyncThriftClient = { + val client = createClient(user, password, host, port, conf) + client.useExistingSessionHandle(sessionHandle) + client + } } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala index d7eab51e818..29676d130de 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/EngineRef.scala @@ -31,7 +31,7 @@ import org.apache.commons.lang3.StringUtils import org.apache.kyuubi.{KYUUBI_VERSION, KyuubiSQLException, Logging, Utils} import org.apache.kyuubi.config.KyuubiConf import org.apache.kyuubi.config.KyuubiConf._ -import org.apache.kyuubi.config.KyuubiReservedKeys.KYUUBI_ENGINE_SUBMIT_TIME_KEY +import org.apache.kyuubi.config.KyuubiReservedKeys.{KYUUBI_ENGINE_SUBMIT_TIME_KEY, KYUUBI_SERVER_HA_NAMESPACE_KEY} import org.apache.kyuubi.engine.EngineType._ import org.apache.kyuubi.engine.ShareLevel.{CONNECTION, GROUP, SERVER, SERVER_LOCAL, ShareLevel} import org.apache.kyuubi.engine.dataagent.DataAgentProcessBuilder @@ -228,6 +228,7 @@ private[kyuubi] class EngineRef( var engineRef = discoveryClient.getServerHost(engineSpace) if (engineRef.nonEmpty) return engineRef.get + conf.set(KYUUBI_SERVER_HA_NAMESPACE_KEY, serverSpace) conf.set(HA_NAMESPACE, engineSpace) conf.set(HA_ENGINE_REF_ID, engineRefId) val started = System.currentTimeMillis() diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/DataAgentResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/DataAgentResource.scala index 60aea983c00..facd347a525 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/DataAgentResource.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/DataAgentResource.scala @@ -18,6 +18,7 @@ package org.apache.kyuubi.server.api.v1 import java.io.{IOException, OutputStreamWriter} +import java.net.ConnectException import java.nio.charset.StandardCharsets import java.util.concurrent.{CompletableFuture, ExecutionException, ExecutorService, RejectedExecutionException, TimeoutException, TimeUnit} import java.util.concurrent.atomic.AtomicBoolean @@ -38,10 +39,16 @@ import org.apache.kyuubi.{KyuubiSQLException, Logging} import org.apache.kyuubi.client.KyuubiSyncThriftClient import org.apache.kyuubi.client.api.v1.dto.{ApprovalRequest, ChatRequest} import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.config.KyuubiConf.ENGINE_SECURITY_ENABLED +import org.apache.kyuubi.ha.HighAvailabilityConf.HA_NAMESPACE +import org.apache.kyuubi.ha.client.{DataAgentSessionRoute, DiscoveryClientProvider} +import org.apache.kyuubi.ha.client.ServiceDiscovery import org.apache.kyuubi.operation.FetchOrientation import org.apache.kyuubi.server.api.ApiRequestContext +import org.apache.kyuubi.service.authentication.InternalSecurityAccessor import org.apache.kyuubi.session.{KyuubiSessionImpl, SessionHandle} import org.apache.kyuubi.shaded.hive.service.rpc.thrift._ +import org.apache.kyuubi.shaded.thrift.transport.TTransportException import org.apache.kyuubi.util.ThreadUtils @Tag(name = "DataAgent") @@ -50,8 +57,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { import DataAgentResource._ - private def verifySessionOwnership(session: KyuubiSessionImpl): Unit = { - val userName = fe.getSessionUser(Map.empty[String, String]) + private def verifySessionOwnership(session: KyuubiSessionImpl, userName: String): Unit = { if (!fe.isAdministrator(userName) && session.user != userName) { throw new ForbiddenException( s"$userName is not allowed to access session ${session.handle}") @@ -59,7 +65,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } // Keep auth failures as 4xx responses before any SSE bytes are sent. - private def resolveAndAuthorize(sessionHandleStr: String): KyuubiSessionImpl = { + private def parseSessionHandle(sessionHandleStr: String): SessionHandle = { val sessionHandle = try { SessionHandle.fromUUID(sessionHandleStr) @@ -67,15 +73,113 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { case _: IllegalArgumentException => throw new WebApplicationException("invalid sessionHandle", 400) } - val session = - try { - fe.be.sessionManager.getSession(sessionHandle).asInstanceOf[KyuubiSessionImpl] - } catch { - case _: KyuubiSQLException => + sessionHandle + } + + private case class ResolvedClient( + client: KyuubiSyncThriftClient, + localSession: Option[KyuubiSessionImpl]) { + def closeTransport(): Unit = if (localSession.isEmpty) client.closeTransport() + } + + private def routeExists(sessionHandleStr: String): Boolean = + DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => + DataAgentSessionRoute.find( + discovery, + fe.getConf.get(HA_NAMESPACE), + sessionHandleStr).isDefined + } + + private def expireRoute(sessionHandleStr: String): Nothing = { + DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => + DataAgentSessionRoute.unregister( + discovery, + fe.getConf.get(HA_NAMESPACE), + sessionHandleStr) + } + throw new WebApplicationException("session expired", 410) + } + + private def resolveClient(sessionHandleStr: String): ResolvedClient = { + val sessionHandle = parseSessionHandle(sessionHandleStr) + val userName = fe.getSessionUser(Map.empty[String, String]) + fe.be.sessionManager.getSessionOption(sessionHandle) match { + case Some(session: KyuubiSessionImpl) => + verifySessionOwnership(session, userName) + val timeout = fe.getConf.get(KyuubiConf.FRONTEND_DATA_AGENT_OPERATION_TIMEOUT) + session.launchEngineOp.getBackgroundHandle.get(timeout, TimeUnit.MILLISECONDS) + if (session.client == null) { + throw new WebApplicationException("Engine session is not ready", 503) + } + if (ServiceDiscovery.supportServiceDiscovery(fe.getConf) && + !routeExists(sessionHandleStr)) { + try fe.be.closeSession(session.handle) + catch { + case NonFatal(e) => warn(s"Failed to clean up expired session ${session.handle}", e) + } + throw new WebApplicationException("session expired", 410) + } + ResolvedClient(session.client, Some(session)) + case _ => + if (!ServiceDiscovery.supportServiceDiscovery(fe.getConf)) { throw new WebApplicationException("session not found", 404) - } - verifySessionOwnership(session) - session + } + val (route, hostPort) = DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { + discovery => + val serverSpace = fe.getConf.get(HA_NAMESPACE) + val found = DataAgentSessionRoute.find(discovery, serverSpace, sessionHandleStr) + .getOrElse(throw new WebApplicationException("session not found", 404)) + if (!fe.isAdministrator(userName) && found.user != userName) { + throw new WebApplicationException("session not found", 404) + } + val hostPort = discovery.getEngineByRefId(found.engineSpace, found.engineRefId) + .getOrElse { + DataAgentSessionRoute.unregister(discovery, serverSpace, sessionHandleStr) + throw new WebApplicationException("session expired", 410) + } + (found, hostPort) + } + val password = if (fe.getConf.get(ENGINE_SECURITY_ENABLED)) { + InternalSecurityAccessor.get().issueToken() + } else { + "anonymous" + } + val client = + try { + KyuubiSyncThriftClient.createClientWithExistingSessionHandle( + route.user, + password, + hostPort._1, + hostPort._2, + fe.getConf, + sessionHandle) + } catch { + case e: TTransportException if e.getCause.isInstanceOf[ConnectException] => + expireRoute(sessionHandleStr) + } + ResolvedClient(client, None) + } + } + + @GET + @Path("sessions/{sessionHandle}") + @Produces(Array(MediaType.APPLICATION_JSON)) + def getSession(@PathParam("sessionHandle") sessionHandleStr: String): String = { + val resolved = resolveClient(sessionHandleStr) + try { + resolved.client.getInfo(TGetInfoType.CLI_DBMS_VER) + "{\"status\":\"ok\"}" + } finally resolved.closeTransport() + } + + @DELETE + @Path("sessions/{sessionHandle}") + def closeSession(@PathParam("sessionHandle") sessionHandleStr: String): Unit = { + val resolved = resolveClient(sessionHandleStr) + resolved.localSession match { + case Some(session) => fe.be.closeSession(session.handle) + case None => resolved.client.closeSession() + } } @ApiResponse( @@ -117,34 +221,9 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } // Validate and authorize before engine startup; launching can be expensive. - val session = resolveAndAuthorize(sessionHandleStr) + val resolved = resolveClient(sessionHandleStr) + val client = resolved.client val operationTimeoutMs = fe.getConf.get(KyuubiConf.FRONTEND_DATA_AGENT_OPERATION_TIMEOUT) - val client: KyuubiSyncThriftClient = - try { - val launchOp = session.launchEngineOp - try { - launchOp.getBackgroundHandle.get(operationTimeoutMs, TimeUnit.MILLISECONDS) - } catch { - case _: TimeoutException => - sendPreflightSseError(response, "Engine did not start within timeout") - return - case e: ExecutionException => - val errMsg = Option(e.getCause).map(_.getMessage).getOrElse("Engine launch failed") - sendPreflightSseError(response, errMsg) - return - } - val c = session.client - if (c == null) { - sendPreflightSseError(response, "Engine session is not ready after waiting") - return - } - c - } catch { - case NonFatal(e) => - error(s"Error processing chat for session $sessionHandleStr", e) - if (!response.isCommitted) sendPreflightSseError(response, e.getMessage) - return - } val stream = new SseStream(response) val deadlineAt = System.currentTimeMillis() + STREAM_MAX_DURATION_MS @@ -175,6 +254,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { warn(s"Op-submit pool rejected chat for session $sessionHandleStr (queue full)") stream.event("error", buildJsonMessage("Server is busy, please retry")) stream.event("done", "{}") + resolved.closeTransport() return } @@ -220,6 +300,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } } finally { if (opHandle != null) closeOperation(client, opHandle) + resolved.closeTransport() } } @@ -245,11 +326,8 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { !REQUEST_ID_PATTERN.matcher(requestId).matches()) { throw new WebApplicationException("invalid requestId", 400) } - val session = resolveAndAuthorize(sessionHandleStr) - val client = session.client - if (client == null) { - throw new WebApplicationException("Engine session is not ready", 503) - } + val resolved = resolveClient(sessionHandleStr) + val client = resolved.client val statement = if (request.isApproved) { s"__approve:$requestId" @@ -257,12 +335,13 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { s"__deny:$requestId" } - val opHandle = client.executeStatement( - statement, - Map.empty[String, String], - false, - 60000L) + var opHandle: TOperationHandle = null try { + opHandle = client.executeStatement( + statement, + Map.empty[String, String], + false, + 60000L) val rowSet = client.fetchResults(opHandle, FetchOrientation.FETCH_NEXT, 1, false) val rows = extractStringRows(rowSet) rows.headOption.getOrElse { @@ -273,7 +352,8 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { JSON_MAPPER.writeValueAsString(node) } } finally { - closeOperation(client, opHandle) + if (opHandle != null) closeOperation(client, opHandle) + resolved.closeTransport() } } diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceAuthorizationSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceAuthorizationSuite.scala new file mode 100644 index 00000000000..8aa4e35aadd --- /dev/null +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceAuthorizationSuite.scala @@ -0,0 +1,63 @@ +/* + * 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.kyuubi.server.api.v1 + +import java.util.UUID +import javax.ws.rs.core.MediaType + +import org.apache.kyuubi.{KyuubiFunSuite, RestFrontendTestHelper} +import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.config.KyuubiConf.{AUTHENTICATION_CUSTOM_CLASS, AUTHENTICATION_METHOD, SERVER_ADMINISTRATORS} +import org.apache.kyuubi.ha.HighAvailabilityConf.HA_NAMESPACE +import org.apache.kyuubi.ha.client.{DataAgentSessionRoute, DiscoveryClientProvider} +import org.apache.kyuubi.server.http.util.HttpAuthUtils.{basicAuthorizationHeader, AUTHORIZATION_HEADER} +import org.apache.kyuubi.service.authentication.AnonymousAuthenticationProviderImpl + +class DataAgentResourceAuthorizationSuite extends KyuubiFunSuite with RestFrontendTestHelper { + + override protected lazy val conf: KyuubiConf = + KyuubiConf() + .set(AUTHENTICATION_METHOD, Seq("CUSTOM")) + .set(AUTHENTICATION_CUSTOM_CLASS, classOf[AnonymousAuthenticationProviderImpl].getName) + .set(SERVER_ADMINISTRATORS, Set("admin")) + + test("a user cannot resolve or remove another user's route") { + val sessionId = UUID.randomUUID().toString + val route = DataAgentSessionRoute("missing-engine-space", "missing-engine-ref", "owner") + DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + DataAgentSessionRoute.register(discovery, conf.get(HA_NAMESPACE), sessionId, route) + } + + try { + val response = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") + .request(MediaType.APPLICATION_JSON_TYPE) + .header(AUTHORIZATION_HEADER, basicAuthorizationHeader("another-user")) + .get() + assert(response.getStatus === 404) + + val persisted = DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + DataAgentSessionRoute.find(discovery, conf.get(HA_NAMESPACE), sessionId) + } + assert(persisted.contains(route)) + } finally { + DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + DataAgentSessionRoute.unregister(discovery, conf.get(HA_NAMESPACE), sessionId) + } + } + } +} diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceSuite.scala index 740fd0d9779..005034520bd 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceSuite.scala @@ -22,12 +22,14 @@ import java.util.UUID import javax.servlet.{ServletOutputStream, WriteListener} import javax.servlet.http.HttpServletResponse import javax.ws.rs.client.Entity -import javax.ws.rs.core.MediaType +import javax.ws.rs.core.{MediaType, Response} import org.mockito.Mockito.{mock, verify, when} import org.apache.kyuubi.{KyuubiFunSuite, RestFrontendTestHelper} import org.apache.kyuubi.client.api.v1.dto.{ApprovalRequest, ChatRequest} +import org.apache.kyuubi.ha.HighAvailabilityConf.{HA_ENGINE_REF_ID, HA_NAMESPACE} +import org.apache.kyuubi.ha.client.{DataAgentSessionRoute, DiscoveryClientProvider} class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper { @@ -48,6 +50,77 @@ class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper assert(response.getStatus === 404) } + test("get and delete validate sessionHandle") { + Seq("not-a-uuid" -> 400, UUID.randomUUID().toString -> 404).foreach { + case (sessionId, expectedStatus) => + Seq("GET", "DELETE").foreach { method => + val response = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") + .request(MediaType.APPLICATION_JSON_TYPE) + .method(method) + assert(response.getStatus === expectedStatus) + } + } + } + + test("stale routes are removed by every session endpoint") { + val requests: Seq[(String, String => Response)] = Seq( + "get" -> { sessionId => + webTarget.path(s"api/v1/data-agent/sessions/$sessionId").request().get() + }, + "delete" -> { sessionId => + webTarget.path(s"api/v1/data-agent/sessions/$sessionId").request().delete() + }, + "chat" -> { sessionId => + webTarget.path(s"api/v1/data-agent/$sessionId/chat").request() + .post(Entity.json(new ChatRequest("hi"))) + }, + "approve" -> { sessionId => + webTarget.path(s"api/v1/data-agent/$sessionId/approve").request() + .post(Entity.json(new ApprovalRequest("request-id", true))) + }) + + requests.foreach { case (endpoint, request) => + withClue(endpoint) { + val sessionId = UUID.randomUUID().toString + registerRoute(sessionId, "anonymous") + assert(request(sessionId).getStatus === 410) + assert(findRoute(sessionId).isEmpty) + } + } + } + + test("connection refusal removes a stale route and returns 410") { + val sessionId = UUID.randomUUID().toString + val engineSpace = s"/unreachable-engine-$sessionId" + val engineRefId = s"unreachable-ref-$sessionId" + val engineConf = conf.clone.set(HA_ENGINE_REF_ID, engineRefId) + val servicePath = DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + val path = discovery.createAndGetServiceNode( + engineConf, + engineSpace, + "127.0.0.1:1", + external = true) + DataAgentSessionRoute.register( + discovery, + conf.get(HA_NAMESPACE), + sessionId, + DataAgentSessionRoute(engineSpace, engineRefId, "anonymous")) + path + } + + try { + val response = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") + .request(MediaType.APPLICATION_JSON_TYPE) + .get() + assert(response.getStatus === 410) + assert(findRoute(sessionId).isEmpty) + } finally { + DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + if (discovery.pathExists(servicePath)) discovery.delete(servicePath) + } + } + } + // -- approve preflight ------------------------------------------------------ test("approve returns 400 for null body") { @@ -175,6 +248,20 @@ class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper (new SseStream(response), sink) } + private def registerRoute(sessionId: String, user: String): DataAgentSessionRoute = { + val route = DataAgentSessionRoute("missing-engine-space", "missing-engine-ref", user) + DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + DataAgentSessionRoute.register(discovery, conf.get(HA_NAMESPACE), sessionId, route) + } + route + } + + private def findRoute(sessionId: String): Option[DataAgentSessionRoute] = { + DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + DataAgentSessionRoute.find(discovery, conf.get(HA_NAMESPACE), sessionId) + } + } + private def servletStreamOver(sink: OutputStream): ServletOutputStream = new ServletOutputStream { override def write(b: Int): Unit = sink.write(b) diff --git a/kyuubi-server/web-ui/src/api/data-agent/index.ts b/kyuubi-server/web-ui/src/api/data-agent/index.ts index 69be4942dd7..b2c286c36c8 100644 --- a/kyuubi-server/web-ui/src/api/data-agent/index.ts +++ b/kyuubi-server/web-ui/src/api/data-agent/index.ts @@ -25,7 +25,7 @@ export interface SessionOpenRequest { export interface SessionHandle { identifier: string - secretId: string + kyuubiInstance: string } export interface SseEvent { @@ -43,14 +43,14 @@ export function openSession(data: SessionOpenRequest) { export function closeSession(sessionHandle: string) { return request({ - url: `api/v1/sessions/${sessionHandle}`, + url: `api/v1/data-agent/sessions/${sessionHandle}`, method: 'delete' }) } export function getSession(sessionHandle: string) { return request({ - url: `api/v1/sessions/${sessionHandle}`, + url: `api/v1/data-agent/sessions/${sessionHandle}`, method: 'get' }) }