diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala index 8722d5b4efb..cce087248c3 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala @@ -247,9 +247,12 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging { } request.setBatchType(request.getBatchType.toUpperCase(Locale.ROOT)) - val userName = fe.getSessionUser(request.getConf.asScala.toMap) + val requestConf = Metadata.sanitizeRequestConf(request.getConf.asScala.toMap) + request.setConf(requestConf.asJava) + + val userName = fe.getSessionUser(requestConf) val ipAddress = fe.getIpAddress - val userProvidedBatchId = request.getConf.asScala.get(KYUUBI_BATCH_ID_KEY) + val userProvidedBatchId = requestConf.get(KYUUBI_BATCH_ID_KEY) userProvidedBatchId.foreach { batchId => try UUID.fromString(batchId) catch { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala index 0553cf90b54..66332419e95 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala @@ -95,4 +95,14 @@ case class Metadata( } def appState: Option[ApplicationState] = Option(engineState).map(ApplicationState.withName) + + def withSanitizedRequestConf: Metadata = { + copy(requestConf = Metadata.sanitizeRequestConf(requestConf)) + } +} + +object Metadata { + def sanitizeRequestConf(requestConf: Map[String, String]): Map[String, String] = { + Option(requestConf).getOrElse(Map.empty).filter { case (_, value) => value != null } + } } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala index 344da0e71e8..cc53ab68ae7 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala @@ -158,11 +158,13 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { // scalastyle:on val username = Option(user).filter(_.nonEmpty).getOrElse("anonymous") val sessionConf = this.getConf.getUserDefaults(user) + val sanitizedConf = Metadata.sanitizeRequestConf(conf) + val sanitizedMetadata = metadata.map(_.withSanitizedRequestConf) new KyuubiBatchSession( username, password, ipAddress, - conf, + sanitizedConf, this, sessionConf, batchType, @@ -170,7 +172,7 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { resource, className, batchArgs, - metadata, + sanitizedMetadata, fromRecovery) } @@ -242,7 +244,7 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { resource = batchRequest.getResource, className = batchRequest.getClassName, requestName = batchRequest.getName, - requestConf = conf, + requestConf = Metadata.sanitizeRequestConf(conf), requestArgs = batchRequest.getArgs.asScala.toSeq, createTime = System.currentTimeMillis(), engineType = batchRequest.getBatchType, @@ -319,9 +321,7 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { stateToRecover.toString, kyuubiInstance, 0, - Int.MaxValue).map { metadata => - createBatchSessionFromRecovery(metadata) - }).getOrElse(Seq.empty) + Int.MaxValue).map(createBatchSessionFromRecovery)).getOrElse(Seq.empty) } } @@ -338,7 +338,7 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { getBatchMetadata(batchId) .filter(m => m.kyuubiInstance == kyuubiInstance && batchStatesToRecovery.contains(m.opState)) - .flatMap { metadata => Some(createBatchSessionFromRecovery(metadata)) } + .map(createBatchSessionFromRecovery) } } } diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala index 86861b6b9cc..e844c0594e6 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/api/v1/BatchesResourceSuite.scala @@ -260,6 +260,27 @@ abstract class BatchesResourceSuiteBase extends KyuubiFunSuite assert(!deleteBatchResponse.readEntity(classOf[CloseBatchResponse]).isSuccess) } + test("ignore null-valued batch configuration") { + val nullConfigKey = "spark.null.option" + val requestObj = newSparkBatchRequest(Map( + "spark.master" -> "local", + nullConfigKey -> null)) + + val response = webTarget.path("api/v1/batches") + .request(MediaType.APPLICATION_JSON_TYPE) + .header(AUTHORIZATION_HEADER, basicAuthorizationHeader("anonymous")) + .post(Entity.entity(requestObj, MediaType.APPLICATION_JSON_TYPE)) + assert(response.getStatus === 200) + val batch = response.readEntity(classOf[Batch]) + + val sessionManager = fe.be.sessionManager.asInstanceOf[KyuubiSessionManager] + eventually(timeout(10.seconds)) { + val metadata = sessionManager.getBatchMetadata(batch.getId) + assert(metadata.isDefined) + assert(!metadata.get.requestConf.contains(nullConfigKey)) + } + } + test("open batch session with uploading resource") { val requestObj = newSparkBatchRequest(Map("spark.master" -> "local")) val exampleJarFile = Paths.get(sparkBatchTestResource.get).toFile diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/session/KyuubiSessionManagerSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/session/KyuubiSessionManagerSuite.scala new file mode 100644 index 00000000000..9e924368aa2 --- /dev/null +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/session/KyuubiSessionManagerSuite.scala @@ -0,0 +1,147 @@ +/* + * 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.session + +import java.util.UUID + +import org.apache.kyuubi.KyuubiFunSuite +import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.config.KyuubiConf.FRONTEND_PROTOCOLS +import org.apache.kyuubi.config.KyuubiConf.FrontendProtocols +import org.apache.kyuubi.operation.OperationState +import org.apache.kyuubi.server.metadata.MetadataManager +import org.apache.kyuubi.server.metadata.api.Metadata + +class KyuubiSessionManagerSuite extends KyuubiFunSuite { + + test("[KYUUBI #7244] sanitize null-valued batch metadata before submission") { + withSessionManager { sessionManager => + val batchId = UUID.randomUUID().toString + val metadata = newMetadata( + identifier = batchId, + kyuubiInstance = "localhost:10009", + requestConf = Map("spark.master" -> "local", "spark.executor.memory" -> null)) + .copy(state = OperationState.INITIALIZED.toString) + + val session = sessionManager.createBatchSession( + metadata.username, + "anonymous", + metadata.ipAddress, + metadata.requestConf, + metadata.engineType, + Option(metadata.requestName), + metadata.resource, + metadata.className, + metadata.requestArgs, + Some(metadata), + fromRecovery = false) + + assert(session.handle.identifier.toString === batchId) + assert(!session.normalizedConf.contains("spark.executor.memory")) + } + } + + test("[KYUUBI #7244] sanitize null-valued batch metadata during recovery") { + withSessionManager { sessionManager => + val validBatchId = UUID.randomUUID().toString + val invalidBatchId = UUID.randomUUID().toString + val kyuubiInstance = "localhost:10009" + val recoveryMetadata = Seq( + newMetadata( + identifier = validBatchId, + kyuubiInstance = kyuubiInstance, + requestConf = Map("spark.master" -> "local")), + newMetadata( + identifier = invalidBatchId, + kyuubiInstance = kyuubiInstance, + requestConf = Map("spark.master" -> "local", "spark.executor.memory" -> null))) + + sessionManager.metadataManager = Some(new RecoveryMetadataManager(recoveryMetadata)) + + val sessions = sessionManager.getBatchSessionsToRecover(kyuubiInstance) + + assert(sessions.map(_.handle.identifier.toString).toSet === Set(validBatchId, invalidBatchId)) + val recoveredInvalidSession = sessions.find( + _.handle.identifier.toString == invalidBatchId).get + assert(!recoveredInvalidSession.normalizedConf.contains("spark.executor.memory")) + + val recoveredMetadata = sessionManager.getBatchMetadata(invalidBatchId).get + assert(recoveredMetadata.state === OperationState.PENDING.toString) + assert(recoveredMetadata.engineError.isEmpty) + } + } + + private class RecoveryMetadataManager(recoveryMetadata: Seq[Metadata]) extends MetadataManager { + private var updatedMetadata: Option[Metadata] = None + + override def getBatchesRecoveryMetadata( + state: String, + kyuubiInstance: String, + from: Int, + size: Int): Seq[Metadata] = { + recoveryMetadata + .filter(metadata => metadata.state == state && metadata.kyuubiInstance == kyuubiInstance) + .slice(from, from + size) + } + + override def getBatchSessionMetadata(batchId: String): Option[Metadata] = { + updatedMetadata + .filter(_.identifier == batchId) + .orElse(recoveryMetadata.find(_.identifier == batchId)) + } + + override def updateMetadata(metadata: Metadata, asyncRetryOnError: Boolean): Unit = { + updatedMetadata = Some(metadata) + } + } + + private def withSessionManager(f: KyuubiSessionManager => Unit): Unit = { + val sessionManager = new KyuubiSessionManager() + val conf = KyuubiConf() + .set(FRONTEND_PROTOCOLS, Seq(FrontendProtocols.REST.toString)) + try { + sessionManager.initialize(conf) + sessionManager.start() + f(sessionManager) + } finally { + sessionManager.stop() + } + } + + private def newMetadata( + identifier: String, + kyuubiInstance: String, + requestConf: Map[String, String]): Metadata = { + Metadata( + identifier = identifier, + sessionType = SessionType.BATCH, + realUser = "kyuubi", + username = "kyuubi", + ipAddress = "127.0.0.1", + kyuubiInstance = kyuubiInstance, + state = OperationState.PENDING.toString, + resource = "intern", + className = "org.apache.kyuubi.SparkWC", + requestName = "kyuubi_batch", + requestConf = requestConf, + requestArgs = Seq.empty, + createTime = System.currentTimeMillis(), + engineType = "spark", + clusterManager = Some("local")) + } +}