wangzhigang1999 commented on code in PR #7530:
URL: https://github.com/apache/kyuubi/pull/7530#discussion_r3892347961
##########
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:
Addressed with an empty-overlay fast path: the execution path now reuses the
session SQLConf and only clones it when per-statement overrides are present.
##########
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:
Addressed by assigning the merged immutable map back explicitly. This keeps
the accumulated overlay clear and avoids relying on Scala assignment rewriting.
##########
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:
Addressed with the same empty-overlay fast path in plan-only mode, so the
base SQLConf is reused unless an interceptor supplied overrides.
--
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]
