From 5be26e3b4ee17171b41fc2e2cae8ebcb3275c427 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Fri, 10 Jul 2026 21:40:38 +0800 Subject: [PATCH 01/10] [KYUUBI #7529][DATA-AGENT] Support cross-server session routing --- .../DataAgentTBinaryFrontendService.scala | 3 + .../session/DataAgentSessionImpl.scala | 4 + .../session/DataAgentSessionManager.scala | 36 +++- .../kyuubi/config/KyuubiReservedKeys.scala | 1 + .../ha/client/DataAgentSessionRoute.scala | 49 ++++++ .../client/DataAgentSessionRouteSuite.scala | 36 ++++ .../client/KyuubiSyncThriftClient.scala | 24 +++ .../org/apache/kyuubi/engine/EngineRef.scala | 3 +- .../server/api/v1/DataAgentResource.scala | 154 +++++++++++++----- .../web-ui/src/api/data-agent/index.ts | 6 +- 10 files changed, 267 insertions(+), 49 deletions(-) create mode 100644 kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DataAgentSessionRoute.scala create mode 100644 kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala diff --git a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala index 12abf86943f..937144bb26d 100644 --- a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala +++ b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala @@ -34,4 +34,7 @@ class DataAgentTBinaryFrontendService(override val serverable: Serverable) super.attributes ++ conf.getAll .get(KYUUBI_ENGINE_ID).map(id => Map(KYUUBI_ENGINE_ID -> id)).getOrElse(Map.empty) } + + private[dataagent] def engineDiscovery: Option[EngineServiceDiscovery] = + discoveryService.collect { case service: EngineServiceDiscovery => service } } 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..06b7b6d4275 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 @@ -51,6 +51,8 @@ class DataAgentSessionImpl( } throw e } + sessionManager.asInstanceOf[DataAgentSessionManager] + .registerRoute(handle.identifier.toString, user) info(s"The data agent session is started.") } @@ -73,6 +75,8 @@ class DataAgentSessionImpl( override def close(): Unit = { try { + sessionManager.asInstanceOf[DataAgentSessionManager] + .unregisterRoute(handle.identifier.toString) dataAgentProvider.close(handle.identifier.toString) } finally { super.close() 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..b3dd3551411 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, KYUUBI_SESSION_HANDLE_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 import org.apache.kyuubi.operation.OperationManager import org.apache.kyuubi.session.{Session, SessionHandle, SessionManager} import org.apache.kyuubi.shaded.hive.service.rpc.thrift.TProtocolVersion @@ -56,6 +58,36 @@ class DataAgentSessionManager(name: String) } } + private def engineDiscovery = DataAgentEngine.currentEngine + .flatMap(_.frontendServices.collectFirst { + case frontend: DataAgentTBinaryFrontendService => frontend + }).flatMap(_.engineDiscovery) + + 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 <- engineDiscovery + } { + val route = DataAgentSessionRoute(conf.get(HA_NAMESPACE), engineRefId, user) + val path = DataAgentSessionRoute.path(serverSpace, sessionId) + discovery.discoveryClient.create(path, "PERSISTENT") + discovery.discoveryClient.setData(path, DataAgentSessionRoute.encode(route)) + } + } + + private[dataagent] def unregisterRoute(sessionId: String): Unit = { + for { + serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY) + discovery <- engineDiscovery + } { + val path = DataAgentSessionRoute.path(serverSpace, sessionId) + if (discovery.discoveryClient.pathExists(path)) { + discovery.discoveryClient.delete(path) + } + } + } + override def stop(): Unit = { try { dataAgentProvider.stop() 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..39343896783 --- /dev/null +++ b/kyuubi-ha/src/main/scala/org/apache/kyuubi/ha/client/DataAgentSessionRoute.scala @@ -0,0 +1,49 @@ +/* + * 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 + +case class DataAgentSessionRoute(engineSpace: String, engineRefId: String, user: String) + +object DataAgentSessionRoute { + private val ROOT_SUFFIX = "DATA_AGENT_sessions" + + def root(serverSpace: String): String = s"${serverSpace}_$ROOT_SUFFIX" + + def path(serverSpace: String, sessionId: String): String = + DiscoveryPaths.makePath(root(serverSpace), sessionId) + + def encode(route: DataAgentSessionRoute): Array[Byte] = { + Seq(route.engineSpace, route.engineRefId, route.user) + .map(value => + Base64.getUrlEncoder.withoutPadding() + .encodeToString(value.getBytes(StandardCharsets.UTF_8))) + .mkString("\n") + .getBytes(StandardCharsets.UTF_8) + } + + def decode(bytes: Array[Byte]): DataAgentSessionRoute = { + val values = new String(bytes, StandardCharsets.UTF_8).split("\n", -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/DataAgentSessionRouteSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala new file mode 100644 index 00000000000..60b1bd9f79e --- /dev/null +++ b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala @@ -0,0 +1,36 @@ +/* + * 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 org.apache.kyuubi.KyuubiFunSuite + +class DataAgentSessionRouteSuite extends KyuubiFunSuite { + + test("encode and decode route") { + val route = DataAgentSessionRoute( + "kyuubi_1.12.0_USER_DATA_AGENT/alice/ds-1234", + "engine-ref-id", + "alice@example.com") + assert(DataAgentSessionRoute.decode(DataAgentSessionRoute.encode(route)) === route) + } + + test("build route path") { + assert(DataAgentSessionRoute.path("kyuubi", "session-id") === + "/kyuubi_DATA_AGENT_sessions/session-id") + } +} 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..3f0c7d439d0 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 attachSession(sessionHandle: SessionHandle): Unit = { + require(_remoteSessionHandle == null, "A session is already attached") + _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 createAttachedClient( + user: String, + password: String, + host: String, + port: Int, + conf: KyuubiConf, + sessionHandle: SessionHandle): KyuubiSyncThriftClient = { + val client = createClient(user, password, host, port, conf) + client.attachSession(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..d3af2c4c658 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 @@ -38,8 +38,13 @@ 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.util.ThreadUtils @@ -50,8 +55,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 +63,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 +71,95 @@ 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, + session: Option[KyuubiSessionImpl], + closeTransport: () => Unit) + + private def routePath(sessionHandleStr: String): String = + DataAgentSessionRoute.path(fe.getConf.get(HA_NAMESPACE), sessionHandleStr) + + private def routeExists(sessionHandleStr: String): Boolean = + DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => + discovery.pathExists(routePath(sessionHandleStr)) + } + + 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 = DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => + val path = routePath(sessionHandleStr) + if (discovery.pathNonExists(path, isPrefix = false)) { + throw new WebApplicationException("session not found", 404) + } + val found = DataAgentSessionRoute.decode(discovery.getData(path)) + if (!fe.isAdministrator(userName) && found.user != userName) { + throw new WebApplicationException("session not found", 404) + } + val hostPort = discovery.getEngineByRefId(found.engineSpace, found.engineRefId) + .getOrElse { + discovery.delete(path) + throw new WebApplicationException("session expired", 410) + } + (found, hostPort) + } + val password = if (fe.getConf.get(ENGINE_SECURITY_ENABLED)) { + InternalSecurityAccessor.get().issueToken() + } else { + "anonymous" + } + val client = KyuubiSyncThriftClient.createAttachedClient( + route._1.user, + password, + route._2._1, + route._2._2, + fe.getConf, + sessionHandle) + ResolvedClient(client, None, () => client.closeTransport()) + } + } + + @GET + @Path("sessions/{sessionHandle}") + @Produces(Array(MediaType.APPLICATION_JSON)) + def getSession(@PathParam("sessionHandle") sessionHandleStr: String): String = { + val resolved = resolveClient(sessionHandleStr) + try "{\"status\":\"ok\"}" + finally resolved.closeTransport() + } + + @DELETE + @Path("sessions/{sessionHandle}") + def closeSession(@PathParam("sessionHandle") sessionHandleStr: String): Unit = { + val resolved = resolveClient(sessionHandleStr) + resolved.session match { + case Some(session) => fe.be.closeSession(session.handle) + case None => resolved.client.closeSession() + } } @ApiResponse( @@ -117,34 +201,17 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } // Validate and authorize before engine startup; launching can be expensive. - val session = resolveAndAuthorize(sessionHandleStr) - val operationTimeoutMs = fe.getConf.get(KyuubiConf.FRONTEND_DATA_AGENT_OPERATION_TIMEOUT) - val client: KyuubiSyncThriftClient = + val resolved = 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 + resolveClient(sessionHandleStr) } catch { case NonFatal(e) => error(s"Error processing chat for session $sessionHandleStr", e) if (!response.isCommitted) sendPreflightSseError(response, e.getMessage) return } + val client = resolved.client + val operationTimeoutMs = fe.getConf.get(KyuubiConf.FRONTEND_DATA_AGENT_OPERATION_TIMEOUT) val stream = new SseStream(response) val deadlineAt = System.currentTimeMillis() + STREAM_MAX_DURATION_MS @@ -175,6 +242,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 +288,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } } finally { if (opHandle != null) closeOperation(client, opHandle) + resolved.closeTransport() } } @@ -245,11 +314,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 +323,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 +340,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/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' }) } From 2933696f3b20710a63971596c162494753a87af0 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Sat, 11 Jul 2026 14:06:48 +0800 Subject: [PATCH 02/10] [KYUUBI #7529][DATA-AGENT] Harden cross-server session routing --- .../session/DataAgentSessionImpl.scala | 17 +++++-- .../session/DataAgentSessionManager.scala | 9 +--- .../ha/client/DataAgentSessionRoute.scala | 46 +++++++++++++++++++ .../client/DataAgentSessionRouteSuite.scala | 6 +++ .../ha/client/DiscoveryClientTests.scala | 15 ++++++ .../server/api/v1/DataAgentResource.scala | 34 ++++++-------- 6 files changed, 95 insertions(+), 32 deletions(-) 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 06b7b6d4275..dae4f4246d9 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,6 +16,8 @@ */ 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} @@ -75,11 +77,18 @@ class DataAgentSessionImpl( override def close(): Unit = { try { - sessionManager.asInstanceOf[DataAgentSessionManager] - .unregisterRoute(handle.identifier.toString) - dataAgentProvider.close(handle.identifier.toString) + try { + sessionManager.asInstanceOf[DataAgentSessionManager] + .unregisterRoute(handle.identifier.toString) + } catch { + case NonFatal(e) => warn("Failed to unregister data agent session route", e) + } } finally { - super.close() + try { + dataAgentProvider.close(handle.identifier.toString) + } finally { + super.close() + } } } 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 b3dd3551411..c471cb6e6f5 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 @@ -70,9 +70,7 @@ class DataAgentSessionManager(name: String) discovery <- engineDiscovery } { val route = DataAgentSessionRoute(conf.get(HA_NAMESPACE), engineRefId, user) - val path = DataAgentSessionRoute.path(serverSpace, sessionId) - discovery.discoveryClient.create(path, "PERSISTENT") - discovery.discoveryClient.setData(path, DataAgentSessionRoute.encode(route)) + DataAgentSessionRoute.register(discovery.discoveryClient, serverSpace, sessionId, route) } } @@ -81,10 +79,7 @@ class DataAgentSessionManager(name: String) serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY) discovery <- engineDiscovery } { - val path = DataAgentSessionRoute.path(serverSpace, sessionId) - if (discovery.discoveryClient.pathExists(path)) { - discovery.discoveryClient.delete(path) - } + DataAgentSessionRoute.unregister(discovery.discoveryClient, serverSpace, sessionId) } } 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 index 39343896783..3eb00ac87d4 100644 --- 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 @@ -20,6 +20,8 @@ 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 { @@ -30,6 +32,44 @@ object DataAgentSessionRoute { def path(serverSpace: String, sessionId: String): String = DiscoveryPaths.makePath(root(serverSpace), sessionId) + def register( + discovery: DiscoveryClient, + serverSpace: String, + sessionId: String, + route: DataAgentSessionRoute): Unit = { + discovery.create( + DiscoveryPaths.makePath(path(serverSpace, sessionId), encodeEntry(route)), + "PERSISTENT") + } + + def find( + discovery: DiscoveryClient, + serverSpace: String, + sessionId: String): Option[DataAgentSessionRoute] = { + val routePath = path(serverSpace, sessionId) + try { + discovery.getChildren(routePath) match { + case Nil => None + case entry :: Nil => Some(decodeEntry(entry)) + case _ => throw new IllegalStateException(s"Multiple routes found for session $sessionId") + } + } catch { + case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) => None + } + } + + 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) => + } + } + def encode(route: DataAgentSessionRoute): Array[Byte] = { Seq(route.engineSpace, route.engineRefId, route.user) .map(value => @@ -39,6 +79,12 @@ object DataAgentSessionRoute { .getBytes(StandardCharsets.UTF_8) } + private def encodeEntry(route: DataAgentSessionRoute): String = + Base64.getUrlEncoder.withoutPadding().encodeToString(encode(route)) + + private def decodeEntry(entry: String): DataAgentSessionRoute = + decode(Base64.getUrlDecoder.decode(entry)) + def decode(bytes: Array[Byte]): DataAgentSessionRoute = { val values = new String(bytes, StandardCharsets.UTF_8).split("\n", -1) require(values.length == 3, "Invalid Data Agent session route") diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala index 60b1bd9f79e..db5305c56a3 100644 --- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala +++ b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala @@ -33,4 +33,10 @@ class DataAgentSessionRouteSuite extends KyuubiFunSuite { assert(DataAgentSessionRoute.path("kyuubi", "session-id") === "/kyuubi_DATA_AGENT_sessions/session-id") } + + test("reject malformed route") { + intercept[IllegalArgumentException] { + DataAgentSessionRoute.decode("not-a-route".getBytes) + } + } } 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..3e9e7520f19 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,19 @@ trait DiscoveryClientTests extends KyuubiFunSuite { assert(data == dataFromGet) } } + + test("data agent session route lifecycle") { + withDiscoveryClient(conf) { discoveryClient => + val serverSpace = "server-space" + val sessionId = "session-id" + val route = DataAgentSessionRoute("engine-space", "engine-ref-id", "alice@example.com") + + assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty) + DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route) + 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-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 d3af2c4c658..f24a06875f0 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 @@ -79,12 +79,12 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { session: Option[KyuubiSessionImpl], closeTransport: () => Unit) - private def routePath(sessionHandleStr: String): String = - DataAgentSessionRoute.path(fe.getConf.get(HA_NAMESPACE), sessionHandleStr) - private def routeExists(sessionHandleStr: String): Boolean = DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => - discovery.pathExists(routePath(sessionHandleStr)) + DataAgentSessionRoute.find( + discovery, + fe.getConf.get(HA_NAMESPACE), + sessionHandleStr).isDefined } private def resolveClient(sessionHandleStr: String): ResolvedClient = { @@ -112,17 +112,15 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { throw new WebApplicationException("session not found", 404) } val route = DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => - val path = routePath(sessionHandleStr) - if (discovery.pathNonExists(path, isPrefix = false)) { - throw new WebApplicationException("session not found", 404) - } - val found = DataAgentSessionRoute.decode(discovery.getData(path)) + 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 { - discovery.delete(path) + DataAgentSessionRoute.unregister(discovery, serverSpace, sessionHandleStr) throw new WebApplicationException("session expired", 410) } (found, hostPort) @@ -148,8 +146,10 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { @Produces(Array(MediaType.APPLICATION_JSON)) def getSession(@PathParam("sessionHandle") sessionHandleStr: String): String = { val resolved = resolveClient(sessionHandleStr) - try "{\"status\":\"ok\"}" - finally resolved.closeTransport() + try { + resolved.client.getInfo(TGetInfoType.CLI_DBMS_VER) + "{\"status\":\"ok\"}" + } finally resolved.closeTransport() } @DELETE @@ -201,15 +201,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } // Validate and authorize before engine startup; launching can be expensive. - val resolved = - try { - resolveClient(sessionHandleStr) - } catch { - case NonFatal(e) => - error(s"Error processing chat for session $sessionHandleStr", e) - if (!response.isCommitted) sendPreflightSseError(response, e.getMessage) - return - } + val resolved = resolveClient(sessionHandleStr) val client = resolved.client val operationTimeoutMs = fe.getConf.get(KyuubiConf.FRONTEND_DATA_AGENT_OPERATION_TIMEOUT) From b62c40012523a6d0da215aa8e311110b2b914ceb Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Sat, 11 Jul 2026 14:35:21 +0800 Subject: [PATCH 03/10] [KYUUBI #7529][DATA-AGENT] Cover cross-server routing corner cases --- .../session/DataAgentSessionImpl.scala | 15 ++- .../ha/client/DataAgentSessionRoute.scala | 35 ++++--- .../client/DataAgentSessionRouteSuite.scala | 16 ++++ .../ha/client/DiscoveryClientTests.scala | 16 ++++ .../etcd/EtcdDiscoveryClientSuite.scala | 4 + .../server/api/v1/DataAgentResource.scala | 42 +++++---- .../DataAgentResourceAuthorizationSuite.scala | 63 +++++++++++++ .../api/v1/DataAgentResourceSuite.scala | 91 +++++++++++++++++++ 8 files changed, 242 insertions(+), 40 deletions(-) create mode 100644 kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/DataAgentResourceAuthorizationSuite.scala 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 dae4f4246d9..56d16f02aca 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 @@ -20,7 +20,7 @@ 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( @@ -29,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 { @@ -53,8 +53,6 @@ class DataAgentSessionImpl( } throw e } - sessionManager.asInstanceOf[DataAgentSessionManager] - .registerRoute(handle.identifier.toString, user) info(s"The data agent session is started.") } @@ -78,8 +76,7 @@ class DataAgentSessionImpl( override def close(): Unit = { try { try { - sessionManager.asInstanceOf[DataAgentSessionManager] - .unregisterRoute(handle.identifier.toString) + dataAgentSessionManager.unregisterRoute(handle.identifier.toString) } catch { case NonFatal(e) => warn("Failed to unregister data agent session route", e) } 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 index 3eb00ac87d4..5d13b5683b3 100644 --- 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 @@ -37,9 +37,19 @@ object DataAgentSessionRoute { serverSpace: String, sessionId: String, route: DataAgentSessionRoute): Unit = { - discovery.create( - DiscoveryPaths.makePath(path(serverSpace, sessionId), encodeEntry(route)), - "PERSISTENT") + val routePath = path(serverSpace, sessionId) + val entry = encodeEntry(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( @@ -47,14 +57,10 @@ object DataAgentSessionRoute { serverSpace: String, sessionId: String): Option[DataAgentSessionRoute] = { val routePath = path(serverSpace, sessionId) - try { - discovery.getChildren(routePath) match { - case Nil => None - case entry :: Nil => Some(decodeEntry(entry)) - case _ => throw new IllegalStateException(s"Multiple routes found for session $sessionId") - } - } catch { - case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) => None + children(discovery, routePath) match { + case Nil => None + case entry :: Nil => Some(decodeEntry(entry)) + case _ => throw new IllegalStateException(s"Multiple routes found for session $sessionId") } } @@ -85,6 +91,13 @@ object DataAgentSessionRoute { private def decodeEntry(entry: String): DataAgentSessionRoute = decode(Base64.getUrlDecoder.decode(entry)) + private def children(discovery: DiscoveryClient, routePath: String): List[String] = { + try discovery.getChildren(routePath) + catch { + case NonFatal(_) if discovery.pathNonExists(routePath, isPrefix = true) => Nil + } + } + def decode(bytes: Array[Byte]): DataAgentSessionRoute = { val values = new String(bytes, StandardCharsets.UTF_8).split("\n", -1) require(values.length == 3, "Invalid Data Agent session route") diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala index db5305c56a3..d187ce63a24 100644 --- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala +++ b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala @@ -29,6 +29,16 @@ class DataAgentSessionRouteSuite extends KyuubiFunSuite { assert(DataAgentSessionRoute.decode(DataAgentSessionRoute.encode(route)) === route) } + test("encode and decode route with unicode and delimiters") { + val namespace = Seq(0x547D, 0x540D, 0x7A7A, 0x95F4).map(_.toChar).mkString + val user = Seq(0x7528, 0x6237).map(_.toChar).mkString + val route = DataAgentSessionRoute( + s"$namespace/with spaces", + "ref:id/with-delimiters", + s"$user+data@example.com") + assert(DataAgentSessionRoute.decode(DataAgentSessionRoute.encode(route)) === route) + } + test("build route path") { assert(DataAgentSessionRoute.path("kyuubi", "session-id") === "/kyuubi_DATA_AGENT_sessions/session-id") @@ -39,4 +49,10 @@ class DataAgentSessionRouteSuite extends KyuubiFunSuite { DataAgentSessionRoute.decode("not-a-route".getBytes) } } + + test("reject route with malformed base64 value") { + intercept[IllegalArgumentException] { + DataAgentSessionRoute.decode("%%%%\nZW5naW5l\ndXNlcg".getBytes) + } + } } 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 3e9e7520f19..ae515c58f5f 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 @@ -172,6 +172,10 @@ trait DiscoveryClientTests extends KyuubiFunSuite { } test("data agent session route lifecycle") { + testDataAgentSessionRouteLifecycle() + } + + protected def testDataAgentSessionRouteLifecycle(): Unit = { withDiscoveryClient(conf) { discoveryClient => val serverSpace = "server-space" val sessionId = "session-id" @@ -179,7 +183,19 @@ trait DiscoveryClientTests extends KyuubiFunSuite { 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/server/api/v1/DataAgentResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/DataAgentResource.scala index f24a06875f0..38e74418141 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 @@ -76,8 +76,9 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { private case class ResolvedClient( client: KyuubiSyncThriftClient, - session: Option[KyuubiSessionImpl], - closeTransport: () => Unit) + localSession: Option[KyuubiSessionImpl]) { + def closeTransport(): Unit = if (localSession.isEmpty) client.closeTransport() + } private def routeExists(sessionHandleStr: String): Boolean = DiscoveryClientProvider.withDiscoveryClient(fe.getConf) { discovery => @@ -106,24 +107,25 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } throw new WebApplicationException("session expired", 410) } - ResolvedClient(session.client, Some(session), () => ()) + ResolvedClient(session.client, Some(session)) case _ => if (!ServiceDiscovery.supportServiceDiscovery(fe.getConf)) { throw new WebApplicationException("session not found", 404) } - val route = 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) + 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) } - (found, hostPort) + 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() @@ -131,13 +133,13 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { "anonymous" } val client = KyuubiSyncThriftClient.createAttachedClient( - route._1.user, + route.user, password, - route._2._1, - route._2._2, + hostPort._1, + hostPort._2, fe.getConf, sessionHandle) - ResolvedClient(client, None, () => client.closeTransport()) + ResolvedClient(client, None) } } @@ -156,7 +158,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { @Path("sessions/{sessionHandle}") def closeSession(@PathParam("sessionHandle") sessionHandleStr: String): Unit = { val resolved = resolveClient(sessionHandleStr) - resolved.session match { + resolved.localSession match { case Some(session) => fe.be.closeSession(session.handle) case None => resolved.client.closeSession() } 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..03750223816 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 @@ -28,6 +28,8 @@ 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_NAMESPACE +import org.apache.kyuubi.ha.client.{DataAgentSessionRoute, DiscoveryClientProvider} class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper { @@ -48,6 +50,81 @@ class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper assert(response.getStatus === 404) } + test("get and delete return 400 for malformed sessionHandle") { + val getResponse = webTarget.path("api/v1/data-agent/sessions/not-a-uuid") + .request(MediaType.APPLICATION_JSON_TYPE) + .get() + assert(getResponse.getStatus === 400) + + val deleteResponse = webTarget.path("api/v1/data-agent/sessions/not-a-uuid") + .request(MediaType.APPLICATION_JSON_TYPE) + .delete() + assert(deleteResponse.getStatus === 400) + } + + test("get and delete return 404 for unknown sessionHandle") { + val unknown = UUID.randomUUID().toString + val getResponse = webTarget.path(s"api/v1/data-agent/sessions/$unknown") + .request(MediaType.APPLICATION_JSON_TYPE) + .get() + assert(getResponse.getStatus === 404) + + val deleteResponse = webTarget.path(s"api/v1/data-agent/sessions/$unknown") + .request(MediaType.APPLICATION_JSON_TYPE) + .delete() + assert(deleteResponse.getStatus === 404) + } + + test("stale route returns 410 and is removed") { + val sessionId = UUID.randomUUID().toString + registerRoute(sessionId, "anonymous") + + 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) + + val secondResponse = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") + .request(MediaType.APPLICATION_JSON_TYPE) + .get() + assert(secondResponse.getStatus === 404) + } + + test("delete removes a stale route and returns 410") { + val sessionId = UUID.randomUUID().toString + registerRoute(sessionId, "anonymous") + + val response = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") + .request(MediaType.APPLICATION_JSON_TYPE) + .delete() + assert(response.getStatus === 410) + assert(findRoute(sessionId).isEmpty) + } + + test("chat removes a stale route before opening the SSE stream") { + val sessionId = UUID.randomUUID().toString + registerRoute(sessionId, "anonymous") + + val response = webTarget.path(s"api/v1/data-agent/$sessionId/chat") + .request(MediaType.APPLICATION_JSON_TYPE) + .post(Entity.entity(new ChatRequest("hi"), MediaType.APPLICATION_JSON_TYPE)) + assert(response.getStatus === 410) + assert(findRoute(sessionId).isEmpty) + } + + test("approve removes a stale route and returns 410") { + val sessionId = UUID.randomUUID().toString + registerRoute(sessionId, "anonymous") + val request = new ApprovalRequest("request-id", true) + + val response = webTarget.path(s"api/v1/data-agent/$sessionId/approve") + .request(MediaType.APPLICATION_JSON_TYPE) + .post(Entity.entity(request, MediaType.APPLICATION_JSON_TYPE)) + assert(response.getStatus === 410) + assert(findRoute(sessionId).isEmpty) + } + // -- approve preflight ------------------------------------------------------ test("approve returns 400 for null body") { @@ -175,6 +252,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) From 0a8eb4ad2eb3180b422da830449f24a78ff52528 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Sat, 11 Jul 2026 15:08:47 +0800 Subject: [PATCH 04/10] [KYUUBI #7529][DATA-AGENT] Expire unreachable session routes --- .../server/api/v1/DataAgentResource.scala | 32 +++++++++++---- .../api/v1/DataAgentResourceSuite.scala | 39 ++++++++++++++++++- 2 files changed, 63 insertions(+), 8 deletions(-) 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 38e74418141..59db80bb3b9 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 @@ -47,6 +48,7 @@ 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") @@ -88,6 +90,16 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { 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]) @@ -132,13 +144,19 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } else { "anonymous" } - val client = KyuubiSyncThriftClient.createAttachedClient( - route.user, - password, - hostPort._1, - hostPort._2, - fe.getConf, - sessionHandle) + val client = + try { + KyuubiSyncThriftClient.createAttachedClient( + route.user, + password, + hostPort._1, + hostPort._2, + fe.getConf, + sessionHandle) + } catch { + case e: TTransportException if e.getCause.isInstanceOf[ConnectException] => + expireRoute(sessionHandleStr) + } ResolvedClient(client, None) } } 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 03750223816..d75fa79ea8e 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 @@ -28,7 +28,7 @@ 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_NAMESPACE +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 { @@ -125,6 +125,43 @@ class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper 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) + + val secondResponse = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") + .request(MediaType.APPLICATION_JSON_TYPE) + .get() + assert(secondResponse.getStatus === 404) + } finally { + DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => + if (discovery.pathExists(servicePath)) discovery.delete(servicePath) + } + } + } + // -- approve preflight ------------------------------------------------------ test("approve returns 400 for null body") { From 3b94883f6bc2fa077f0a02ebc9da1eb849440839 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Mon, 13 Jul 2026 10:49:39 +0800 Subject: [PATCH 05/10] [KYUUBI #7529][DATA-AGENT] Simplify session route encoding --- .../ha/client/DataAgentSessionRoute.scala | 25 ++++---- .../client/DataAgentSessionRouteSuite.scala | 58 ------------------- .../ha/client/DiscoveryClientTests.scala | 5 +- 3 files changed, 14 insertions(+), 74 deletions(-) delete mode 100644 kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala 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 index 5d13b5683b3..cedd35868b8 100644 --- 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 @@ -38,7 +38,7 @@ object DataAgentSessionRoute { sessionId: String, route: DataAgentSessionRoute): Unit = { val routePath = path(serverSpace, sessionId) - val entry = encodeEntry(route) + val entry = encode(route) children(discovery, routePath) match { case Nil => val entryPath = DiscoveryPaths.makePath(routePath, entry) @@ -59,7 +59,7 @@ object DataAgentSessionRoute { val routePath = path(serverSpace, sessionId) children(discovery, routePath) match { case Nil => None - case entry :: Nil => Some(decodeEntry(entry)) + case entry :: Nil => Some(decode(entry)) case _ => throw new IllegalStateException(s"Multiple routes found for session $sessionId") } } @@ -76,20 +76,15 @@ object DataAgentSessionRoute { } } - def encode(route: DataAgentSessionRoute): Array[Byte] = { + private def encode(route: DataAgentSessionRoute): String = { Seq(route.engineSpace, route.engineRefId, route.user) - .map(value => - Base64.getUrlEncoder.withoutPadding() - .encodeToString(value.getBytes(StandardCharsets.UTF_8))) - .mkString("\n") - .getBytes(StandardCharsets.UTF_8) + .map(encodeValue) + .mkString(".") } - private def encodeEntry(route: DataAgentSessionRoute): String = - Base64.getUrlEncoder.withoutPadding().encodeToString(encode(route)) - - private def decodeEntry(entry: String): DataAgentSessionRoute = - decode(Base64.getUrlDecoder.decode(entry)) + 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) @@ -98,8 +93,8 @@ object DataAgentSessionRoute { } } - def decode(bytes: Array[Byte]): DataAgentSessionRoute = { - val values = new String(bytes, StandardCharsets.UTF_8).split("\n", -1) + 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) diff --git a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala b/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala deleted file mode 100644 index d187ce63a24..00000000000 --- a/kyuubi-ha/src/test/scala/org/apache/kyuubi/ha/client/DataAgentSessionRouteSuite.scala +++ /dev/null @@ -1,58 +0,0 @@ -/* - * 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 org.apache.kyuubi.KyuubiFunSuite - -class DataAgentSessionRouteSuite extends KyuubiFunSuite { - - test("encode and decode route") { - val route = DataAgentSessionRoute( - "kyuubi_1.12.0_USER_DATA_AGENT/alice/ds-1234", - "engine-ref-id", - "alice@example.com") - assert(DataAgentSessionRoute.decode(DataAgentSessionRoute.encode(route)) === route) - } - - test("encode and decode route with unicode and delimiters") { - val namespace = Seq(0x547D, 0x540D, 0x7A7A, 0x95F4).map(_.toChar).mkString - val user = Seq(0x7528, 0x6237).map(_.toChar).mkString - val route = DataAgentSessionRoute( - s"$namespace/with spaces", - "ref:id/with-delimiters", - s"$user+data@example.com") - assert(DataAgentSessionRoute.decode(DataAgentSessionRoute.encode(route)) === route) - } - - test("build route path") { - assert(DataAgentSessionRoute.path("kyuubi", "session-id") === - "/kyuubi_DATA_AGENT_sessions/session-id") - } - - test("reject malformed route") { - intercept[IllegalArgumentException] { - DataAgentSessionRoute.decode("not-a-route".getBytes) - } - } - - test("reject route with malformed base64 value") { - intercept[IllegalArgumentException] { - DataAgentSessionRoute.decode("%%%%\nZW5naW5l\ndXNlcg".getBytes) - } - } -} 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 ae515c58f5f..8313d626230 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 @@ -179,7 +179,10 @@ trait DiscoveryClientTests extends KyuubiFunSuite { withDiscoveryClient(conf) { discoveryClient => val serverSpace = "server-space" val sessionId = "session-id" - val route = DataAgentSessionRoute("engine-space", "engine-ref-id", "alice@example.com") + val route = DataAgentSessionRoute( + "engine space/with.delimiters", + "engine-ref:id/with.delimiters", + "alice+测试@example.com") assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty) DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route) From 872d8198e38b6c38ae061368ce99f16a055a3322 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Mon, 13 Jul 2026 10:52:09 +0800 Subject: [PATCH 06/10] [KYUUBI #7529][DATA-AGENT] Trim session route lifecycle changes --- .../session/DataAgentSessionImpl.scala | 17 +++++++---------- .../session/DataAgentSessionManager.scala | 7 ++----- 2 files changed, 9 insertions(+), 15 deletions(-) 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 56d16f02aca..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 @@ -75,17 +75,14 @@ class DataAgentSessionImpl( override def close(): Unit = { try { - try { - dataAgentSessionManager.unregisterRoute(handle.identifier.toString) - } catch { - case NonFatal(e) => warn("Failed to unregister data agent session route", e) - } + 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 { - try { - dataAgentProvider.close(handle.identifier.toString) - } finally { - super.close() - } + super.close() } } 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 c471cb6e6f5..c2a708d51bb 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,7 +18,7 @@ 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_SERVER_HA_NAMESPACE_KEY, 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, DataAgentTBinaryFrontendService} import org.apache.kyuubi.engine.dataagent.operation.DataAgentOperationManager @@ -52,10 +52,7 @@ 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 engineDiscovery = DataAgentEngine.currentEngine From 4dfe800cb08c543a96253d30b32b4095b037a688 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Mon, 13 Jul 2026 11:01:42 +0800 Subject: [PATCH 07/10] [KYUUBI #7529][DATA-AGENT] Consolidate route endpoint tests --- .../api/v1/DataAgentResourceSuite.scala | 113 ++++++------------ 1 file changed, 36 insertions(+), 77 deletions(-) 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 d75fa79ea8e..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,7 +22,7 @@ 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} @@ -50,79 +50,43 @@ class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper assert(response.getStatus === 404) } - test("get and delete return 400 for malformed sessionHandle") { - val getResponse = webTarget.path("api/v1/data-agent/sessions/not-a-uuid") - .request(MediaType.APPLICATION_JSON_TYPE) - .get() - assert(getResponse.getStatus === 400) - - val deleteResponse = webTarget.path("api/v1/data-agent/sessions/not-a-uuid") - .request(MediaType.APPLICATION_JSON_TYPE) - .delete() - assert(deleteResponse.getStatus === 400) - } - - test("get and delete return 404 for unknown sessionHandle") { - val unknown = UUID.randomUUID().toString - val getResponse = webTarget.path(s"api/v1/data-agent/sessions/$unknown") - .request(MediaType.APPLICATION_JSON_TYPE) - .get() - assert(getResponse.getStatus === 404) - - val deleteResponse = webTarget.path(s"api/v1/data-agent/sessions/$unknown") - .request(MediaType.APPLICATION_JSON_TYPE) - .delete() - assert(deleteResponse.getStatus === 404) - } - - test("stale route returns 410 and is removed") { - val sessionId = UUID.randomUUID().toString - registerRoute(sessionId, "anonymous") - - 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) - - val secondResponse = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") - .request(MediaType.APPLICATION_JSON_TYPE) - .get() - assert(secondResponse.getStatus === 404) - } - - test("delete removes a stale route and returns 410") { - val sessionId = UUID.randomUUID().toString - registerRoute(sessionId, "anonymous") - - val response = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") - .request(MediaType.APPLICATION_JSON_TYPE) - .delete() - assert(response.getStatus === 410) - assert(findRoute(sessionId).isEmpty) - } - - test("chat removes a stale route before opening the SSE stream") { - val sessionId = UUID.randomUUID().toString - registerRoute(sessionId, "anonymous") - - val response = webTarget.path(s"api/v1/data-agent/$sessionId/chat") - .request(MediaType.APPLICATION_JSON_TYPE) - .post(Entity.entity(new ChatRequest("hi"), MediaType.APPLICATION_JSON_TYPE)) - assert(response.getStatus === 410) - assert(findRoute(sessionId).isEmpty) + 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("approve removes a stale route and returns 410") { - val sessionId = UUID.randomUUID().toString - registerRoute(sessionId, "anonymous") - val request = new ApprovalRequest("request-id", true) - - val response = webTarget.path(s"api/v1/data-agent/$sessionId/approve") - .request(MediaType.APPLICATION_JSON_TYPE) - .post(Entity.entity(request, MediaType.APPLICATION_JSON_TYPE)) - assert(response.getStatus === 410) - assert(findRoute(sessionId).isEmpty) + 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") { @@ -150,11 +114,6 @@ class DataAgentResourceSuite extends KyuubiFunSuite with RestFrontendTestHelper .get() assert(response.getStatus === 410) assert(findRoute(sessionId).isEmpty) - - val secondResponse = webTarget.path(s"api/v1/data-agent/sessions/$sessionId") - .request(MediaType.APPLICATION_JSON_TYPE) - .get() - assert(secondResponse.getStatus === 404) } finally { DiscoveryClientProvider.withDiscoveryClient(conf) { discovery => if (discovery.pathExists(servicePath)) discovery.delete(servicePath) From ac163773308e1ef62a0fa2e858c3601cb2c194ad Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Mon, 13 Jul 2026 11:04:53 +0800 Subject: [PATCH 08/10] [KYUUBI #7529][DATA-AGENT] Reuse engine service discovery --- .../DataAgentTBinaryFrontendService.scala | 3 --- .../session/DataAgentSessionManager.scala | 16 +++++++++------- .../kyuubi/ha/client/DataAgentSessionRoute.scala | 4 ++-- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala index 937144bb26d..12abf86943f 100644 --- a/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala +++ b/externals/kyuubi-data-agent-engine/src/main/scala/org/apache/kyuubi/engine/dataagent/DataAgentTBinaryFrontendService.scala @@ -34,7 +34,4 @@ class DataAgentTBinaryFrontendService(override val serverable: Serverable) super.attributes ++ conf.getAll .get(KYUUBI_ENGINE_ID).map(id => Map(KYUUBI_ENGINE_ID -> id)).getOrElse(Map.empty) } - - private[dataagent] def engineDiscovery: Option[EngineServiceDiscovery] = - discoveryService.collect { case service: EngineServiceDiscovery => service } } 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 c2a708d51bb..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 @@ -24,7 +24,7 @@ import org.apache.kyuubi.engine.dataagent.{DataAgentEngine, DataAgentTBinaryFron 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 +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 @@ -55,28 +55,30 @@ class DataAgentSessionManager(name: String) new DataAgentSessionImpl(protocol, user, password, ipAddress, conf, this) } - private def engineDiscovery = DataAgentEngine.currentEngine + private def discoveryClient = DataAgentEngine.currentEngine .flatMap(_.frontendServices.collectFirst { case frontend: DataAgentTBinaryFrontendService => frontend - }).flatMap(_.engineDiscovery) + }).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 <- engineDiscovery + discovery <- discoveryClient } { val route = DataAgentSessionRoute(conf.get(HA_NAMESPACE), engineRefId, user) - DataAgentSessionRoute.register(discovery.discoveryClient, serverSpace, sessionId, route) + DataAgentSessionRoute.register(discovery, serverSpace, sessionId, route) } } private[dataagent] def unregisterRoute(sessionId: String): Unit = { for { serverSpace <- conf.getOption(KYUUBI_SERVER_HA_NAMESPACE_KEY) - discovery <- engineDiscovery + discovery <- discoveryClient } { - DataAgentSessionRoute.unregister(discovery.discoveryClient, serverSpace, sessionId) + DataAgentSessionRoute.unregister(discovery, serverSpace, sessionId) } } 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 index cedd35868b8..c07a5d6e7d8 100644 --- 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 @@ -27,9 +27,9 @@ case class DataAgentSessionRoute(engineSpace: String, engineRefId: String, user: object DataAgentSessionRoute { private val ROOT_SUFFIX = "DATA_AGENT_sessions" - def root(serverSpace: String): String = s"${serverSpace}_$ROOT_SUFFIX" + private def root(serverSpace: String): String = s"${serverSpace}_$ROOT_SUFFIX" - def path(serverSpace: String, sessionId: String): String = + private def path(serverSpace: String, sessionId: String): String = DiscoveryPaths.makePath(root(serverSpace), sessionId) def register( From 327202f40bb2f346e74688f5cf07be31114450b3 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Mon, 13 Jul 2026 11:46:58 +0800 Subject: [PATCH 09/10] [KYUUBI #7529][DATA-AGENT] Clarify existing session handle reuse --- .../org/apache/kyuubi/client/KyuubiSyncThriftClient.scala | 8 ++++---- .../apache/kyuubi/server/api/v1/DataAgentResource.scala | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) 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 3f0c7d439d0..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,8 +62,8 @@ class KyuubiSyncThriftClient private ( // Visible for testing. private[kyuubi] def remoteSessionHandle: TSessionHandle = _remoteSessionHandle - private[kyuubi] def attachSession(sessionHandle: SessionHandle): Unit = { - require(_remoteSessionHandle == null, "A session is already attached") + private[kyuubi] def useExistingSessionHandle(sessionHandle: SessionHandle): Unit = { + require(_remoteSessionHandle == null, "A session handle is already set") _remoteSessionHandle = sessionHandle.toTSessionHandle } @@ -527,7 +527,7 @@ private[kyuubi] object KyuubiSyncThriftClient extends Logging { aliveTimeout) } - def createAttachedClient( + def createClientWithExistingSessionHandle( user: String, password: String, host: String, @@ -535,7 +535,7 @@ private[kyuubi] object KyuubiSyncThriftClient extends Logging { conf: KyuubiConf, sessionHandle: SessionHandle): KyuubiSyncThriftClient = { val client = createClient(user, password, host, port, conf) - client.attachSession(sessionHandle) + client.useExistingSessionHandle(sessionHandle) client } } 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 59db80bb3b9..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 @@ -146,7 +146,7 @@ private[v1] class DataAgentResource extends ApiRequestContext with Logging { } val client = try { - KyuubiSyncThriftClient.createAttachedClient( + KyuubiSyncThriftClient.createClientWithExistingSessionHandle( route.user, password, hostPort._1, From b73b0751d212b6a9e9a05a179611337fd4787819 Mon Sep 17 00:00:00 2001 From: wangzhigang Date: Mon, 13 Jul 2026 11:56:10 +0800 Subject: [PATCH 10/10] [KYUUBI #7529][DATA-AGENT] Fix scalastyle violation --- .../org/apache/kyuubi/ha/client/DiscoveryClientTests.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8313d626230..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 @@ -182,7 +182,7 @@ trait DiscoveryClientTests extends KyuubiFunSuite { val route = DataAgentSessionRoute( "engine space/with.delimiters", "engine-ref:id/with.delimiters", - "alice+测试@example.com") + "alice+route@example.com") assert(DataAgentSessionRoute.find(discoveryClient, serverSpace, sessionId).isEmpty) DataAgentSessionRoute.register(discoveryClient, serverSpace, sessionId, route)