RamonZhou commented on code in PR #58264: URL: https://github.com/apache/spark/pull/58264#discussion_r3866685937
########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/PythonWorkerEnvironment.scala: ########## @@ -0,0 +1,219 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connect.service + +import java.nio.charset.StandardCharsets + +import org.apache.spark.{SparkEnv, SparkException} +import org.apache.spark.sql.RuntimeConfig +import org.apache.spark.sql.connect.config.Connect +import org.apache.spark.sql.internal.SQLConf + +/** + * The environment variables that Python worker processes launched for a session's Python + * functions should inherit. + * + * The environment is carried by session configurations under a reserved prefix, one configuration + * per variable: `spark.pythonWorkerEnv.FOO=bar` makes `FOO` visible as `bar` in `os.environ` + * inside a Python UDF. The configurations are the authoritative session state -- no second copy + * of the environment is maintained as session state -- so the environment follows the session + * wherever ordinary session configurations follow it, including into a session created by + * `cloneSession`. A request's snapshot is also held in the plan cache keys of the plans it + * caches, since a cached plan is only reusable by a request carrying the same environment. + * + * Names are preserved case-sensitively by Spark. On a case-sensitive operating system `FOO` and + * `foo` are therefore distinct variables; Windows process environments are case-insensitive, so + * what a worker observes there is the platform's business rather than Spark's. + * + * A request reads the environment once and uses that one snapshot for everything it does, because + * the configurations can change underneath it: another request may set them while this one is + * still planning. Re-reading would let a plan be built with one environment and cached under + * another. + */ +private[connect] object PythonWorkerEnvironment { + + /** Prefix of the session configurations that carry the environment. */ + val confPrefix: String = "spark.pythonWorkerEnv." + + /** + * Environment variable names accepted under [[confPrefix]]. + * + * This is deliberately stricter than the operating system requires. A POSIX environment permits + * any byte except `=` and NUL in a name, and container platforms accept their own broader sets, + * but a name outside this pattern cannot be referenced portably from a shell, so accepting one + * would let a session install a variable that some consumers can never read. It is a + * portability policy, not a description of what a process environment can hold. + */ + val namePattern: String = "^[A-Za-z_][A-Za-z0-9_]*$" + + private val compiledNamePattern = namePattern.r + + // A rejected name can be arbitrarily long, so messages carry a bounded prefix of it rather than + // the whole name. + private val maxNameCharsInMessage = 32 + + /** + * The environment carried by `conf`, without validation. + * + * Callers take one snapshot per request and pass it around. Validation is separate so that the + * plan cache can tell two environments apart without rejecting an invalid one: an invalid entry + * has to fail the queries that would install it in a worker, not every query in the session. + */ + def read(conf: SQLConf): Map[String, String] = extract(conf.getAllConfs) + + /** The environment carried by the configurations in `allConfs`. */ + private def extract(allConfs: Map[String, String]): Map[String, String] = { + allConfs.iterator + .filter { case (key, _) => key.startsWith(confPrefix) } + .map { case (key, value) => key.substring(confPrefix.length) -> value } + .toMap + } + + /** + * Rejects a malformed or oversized environment. + * + * This runs when a Python function is built, which is the one point that every way of writing a + * configuration reaches: the Spark Connect config RPC, SQL `SET`, and the application-level + * configurations merged into a new session all arrive here. [[validateConfigChange]] rejects a + * write through the config RPC earlier and more helpfully, but it cannot see the other two, so + * this is the check that makes an invalid environment unable to reach a worker at all. + * + * A message may name a variable but never carries its value, so a rejection cannot copy a value + * into a log or a stack trace. Note that the name is chosen by the user, so a name is only as + * safe as what the user put in it. + * + * @throws SparkException + * if a name is malformed or too long, a value cannot be carried by a process environment, or + * the collection exceeds a limit. + */ + def validate(variables: Map[String, String]): Unit = { + val conf = SparkEnv.get.conf + val maxCount = conf.get(Connect.CONNECT_PYTHON_WORKER_ENV_MAX_VARIABLES) + val maxNameLength = conf.get(Connect.CONNECT_PYTHON_WORKER_ENV_MAX_NAME_LENGTH) + val maxTotalSizeBytes = conf.get(Connect.CONNECT_PYTHON_WORKER_ENV_MAX_TOTAL_SIZE_BYTES) + + if (variables.size > maxCount) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.PYTHON_WORKER_ENV_TOO_MANY_VARIABLES", + messageParameters = Map( + "count" -> variables.size.toString, + "prefix" -> confPrefix, + "maxCount" -> maxCount.toString), + cause = null) + } + + var totalSizeBytes = 0L + variables.foreach { case (name, value) => + // `matches` requires the whole name to match. Searching for the pattern instead would accept + // a name with a trailing newline, because `$` also matches before a terminating line break. + if (name.length > maxNameLength || !compiledNamePattern.matches(name)) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.INVALID_PYTHON_WORKER_ENV_VAR_NAME", + messageParameters = Map( + "name" -> describeName(name), + "prefix" -> confPrefix, + "pattern" -> namePattern, + "maxLength" -> maxNameLength.toString), + cause = null) + } + // A process environment cannot carry NUL. Rejecting it here rather than letting the worker + // launch fail matters for more than the error message: the launch failure is an + // `IllegalArgumentException` from the JDK whose own message embeds the offending value. + if (value.indexOf(0) >= 0) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.INVALID_PYTHON_WORKER_ENV_VAR_VALUE", + messageParameters = Map("name" -> describeName(name), "prefix" -> confPrefix), + cause = null) + } + totalSizeBytes += utf8Length(name) + utf8Length(value) + } + + if (totalSizeBytes > maxTotalSizeBytes) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.PYTHON_WORKER_ENV_TOO_LARGE", + messageParameters = Map( + "prefix" -> confPrefix, + "size" -> totalSizeBytes.toString, + "maxSize" -> maxTotalSizeBytes.toString), + cause = null) + } + } + + /** + * Rejects a configuration write that would leave the session with an invalid environment. + * + * A no-op for a key outside [[confPrefix]]. For a key under it, the environment that the write + * would produce is validated before the write happens, so an invalid environment never enters + * the session at all and the failure points at the call that caused it. + * + * This covers the Spark Connect config RPC, which is both how a client sets a configuration + * explicitly and how `SparkSession.builder.config` applies one. It does not cover SQL `SET` or + * the application-level configurations merged into a new session: both reach the session + * configurations without passing through the RPC, which is why [[validate]] at build time stays + * as the check that no invalid environment can reach a worker. + * + * Removing a variable is deliberately not validated. A removal can only shrink the environment, + * and it is how a session recovers from an environment that one of those unchecked paths left + * invalid; validating a removal would leave such a session with no way back. + * + * @throws SparkException + * if the environment the write would produce is malformed or oversized. + */ + def validateConfigChange(conf: RuntimeConfig, key: String, value: Option[String]): Unit = { + if (key.startsWith(confPrefix)) { + // An absent value is rejected by `SQLConf` itself. Leave that failure where it is instead of + // reporting a missing value as an invalid environment. + value.foreach { newValue => + validate(extract(conf.getAll) + (key.substring(confPrefix.length) -> newValue)) Review Comment: Fixed by adding `sqlConf.settings.synchronized` to atomize the operation and also tests -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
