Copilot commented on code in PR #7530:
URL: https://github.com/apache/kyuubi/pull/7530#discussion_r3889680479
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala:
##########
@@ -77,7 +78,11 @@ class PlanOnlyStatement(
override protected def runInternal(): Unit =
try {
withLocalProperties {
- SQLConf.withExistingConf(spark.sessionState.conf) {
+ val operationConf = spark.sessionState.conf.clone()
+ confOverlay.foreach { case (key, value) =>
+ operationConf.setConfString(key, value)
+ }
+ SQLConf.withExistingConf(operationConf) {
val parsed = spark.sessionState.sqlParser.parsePlan(statement)
Review Comment:
This clones `spark.sessionState.conf` on every invocation even when
`confOverlay` is empty. Since plan-only mode can be used frequently, avoid the
extra clone by reusing the base `SQLConf` when there are no per-statement
overrides.
##########
kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/StatementInterception.scala:
##########
@@ -0,0 +1,119 @@
+/*
+ * 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.plugin
+
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
+
+import org.apache.kyuubi.KyuubiSQLException
+
+/** A concrete, immutable [[StatementInterceptContext]] built per interceptor
invocation. */
+private[kyuubi] class StatementInterceptContextImpl(
+ override val sessionId: String,
+ override val statementId: String,
+ override val user: String,
+ override val realUser: String,
+ override val ipAddress: String,
+ override val statement: String,
+ override val confOverlay: java.util.Map[String, String],
+ override val runAsync: Boolean,
+ override val queryTimeout: Long,
+ override val engineType: String) extends StatementInterceptContext
+
+private[kyuubi] case class InterceptedStatement(
+ statement: String,
+ confOverlay: Map[String, String])
+
+private[kyuubi] object StatementInterception {
+
+ /**
+ * Initialize interceptors in order. If an interceptor's initialize throws,
every interceptor
+ * whose initialization was attempted is closed in reverse order before the
error is rethrown, so
+ * that a failed server startup does not leak resources (threads,
connections, external clients)
+ * allocated before the failure. Errors raised while closing are attached as
suppressed
+ * exceptions.
+ */
+ def initialize(
+ interceptors: Seq[StatementInterceptor],
+ conf: java.util.Map[String, String]): Unit = {
+ val initialized = new
scala.collection.mutable.ArrayBuffer[StatementInterceptor]()
+ try {
+ interceptors.foreach { interceptor =>
+ initialized += interceptor
+ interceptor.initialize(conf)
+ }
+ } catch {
+ case initError: Throwable =>
+ initialized.reverseIterator.foreach { interceptor =>
+ try {
+ interceptor.close()
+ } catch {
+ case NonFatal(closeError) => initError.addSuppressed(closeError)
+ }
+ }
+ throw initError
+ }
+ }
+
+ /**
+ * Run the interceptor chain over a statement. Each interceptor sees the
statement and config
+ * overlay produced by the previous one; a reject, an exception, or a null
result fails the
+ * statement.
+ *
+ * @param interceptors the interceptors in execution order
+ * @param initialStatement the original statement
+ * @param initialConfOverlay the client-provided per-statement config overlay
+ * @param contextFor builds a context for the current statement and config
overlay
+ * @return the final statement and config overlay to route to the engine
+ */
+ def run(
+ interceptors: Seq[StatementInterceptor],
+ initialStatement: String,
+ initialConfOverlay: Map[String, String],
+ contextFor: (
+ String,
+ java.util.Map[String, String]) => StatementInterceptContext):
InterceptedStatement = {
+ var currentStatement = initialStatement
+ var currentConfOverlay = initialConfOverlay
+ interceptors.foreach { interceptor =>
+ val interceptorName = interceptor.getClass.getName
+ val result =
+ try {
+ interceptor.beforeExecuteStatement(
+ contextFor(currentStatement, currentConfOverlay.asJava))
+ } catch {
+ case e: KyuubiSQLException => throw e
+ case NonFatal(e) =>
+ throw KyuubiSQLException(s"Statement interceptor $interceptorName
failed", e)
+ }
+ if (result == null) {
+ throw KyuubiSQLException(s"Statement interceptor $interceptorName
returned null")
+ }
+ result.action() match {
+ case StatementInterceptResult.Action.PROCEED => // keep the current
statement
+ case StatementInterceptResult.Action.REWRITE => currentStatement =
result.statement()
+ case StatementInterceptResult.Action.REJECT =>
+ // SQLState 42501 = "Insufficient Privilege". Use the specific
access-rule subclass so
+ // clients can distinguish a policy rejection from the generic
syntax-error class 42000.
+ throw KyuubiSQLException(result.message(), sqlState = "42501")
+ }
+ currentConfOverlay ++= result.confOverlay().asScala
Review Comment:
`currentConfOverlay` is typed as an immutable `Map[String, String]`, so
using `++=` here is not valid and will not compile (the immutable collections
API provides `++`, not `++=`). Assign the merged map back to
`currentConfOverlay` instead so interceptor config deltas are accumulated
correctly.
##########
externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala:
##########
@@ -83,14 +85,20 @@ class ExecuteStatement(
protected def executeStatement(): Unit =
try {
withLocalProperties {
- setState(OperationState.RUNNING)
- info(diagnostics)
-
Thread.currentThread().setContextClassLoader(spark.sharedState.jarClassLoader)
- addOperationListener()
- result = spark.sql(statement)
- iter = collectAsIterator(result)
- setCompiledStateIfNeeded()
- setState(OperationState.FINISHED)
+ val operationConf = spark.sessionState.conf.clone()
+ confOverlay.foreach { case (key, value) =>
+ operationConf.setConfString(key, value)
+ }
+ SQLConf.withExistingConf(operationConf) {
Review Comment:
This clones `spark.sessionState.conf` for every statement, even when
`confOverlay` is empty (the common case). Cloning `SQLConf` on every execution
can add avoidable overhead; consider only cloning when there are actual
overrides, and otherwise reuse `spark.sessionState.conf`.
--
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]
