diff --git a/docs/configuration/settings.md b/docs/configuration/settings.md index 020cb0efdd5..8b89d3654c6 100644 --- a/docs/configuration/settings.md +++ b/docs/configuration/settings.md @@ -453,6 +453,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co | kyuubi.operation.result.saveToFile.minSize | 209715200 | The minSize of Spark result save to file, default value is 200 MB.we use spark's `EstimationUtils#getSizePerRowestimate` to estimate the output size of the execution plan. | long | 1.9.0 | | kyuubi.operation.scheduler.pool | <undefined> | The scheduler pool of job. Note that, this config should be used after changing Spark config spark.scheduler.mode=FAIR. | string | 1.1.1 | | kyuubi.operation.spark.listener.enabled | true | When set to true, Spark engine registers an SQLOperationListener before executing the statement, logging a few summary statistics when each stage completes. | boolean | 1.6.0 | +| kyuubi.operation.statement.interceptors | <undefined> | A comma-separated list of statement interceptor plugins for Kyuubi Server. Each value should be a subclass of `org.apache.kyuubi.plugin.StatementInterceptor` with a zero-arg constructor. They are invoked in the configured order on the server before each interactive statement is routed to the engine, and can inspect, reject, or rewrite the statement. | seq | 1.12.0 | | kyuubi.operation.status.polling.timeout | PT5S | Timeout(ms) for long polling asynchronous running sql query's status | duration | 1.0.0 | | kyuubi.operation.timeout.pool.keepalive.time | PT1M | Keep-alive time for idle threads in the timeout scheduler pool. | duration | 1.11.0 | | kyuubi.operation.timeout.pool.size | 8 | Number of threads in the timeout scheduler pool used for operation timeout monitoring. | int | 1.11.0 | diff --git a/docs/extensions/server/index.rst b/docs/extensions/server/index.rst index d8c860d6b4c..0558397f5c1 100644 --- a/docs/extensions/server/index.rst +++ b/docs/extensions/server/index.rst @@ -25,5 +25,6 @@ ability of kyuubi servers. authentication configuration + statement_interceptor events applications diff --git a/docs/extensions/server/statement_interceptor.rst b/docs/extensions/server/statement_interceptor.rst new file mode 100644 index 00000000000..cefe9618225 --- /dev/null +++ b/docs/extensions/server/statement_interceptor.rst @@ -0,0 +1,174 @@ +.. 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. + +Intercept Statements with Custom Statement Interceptor +====================================================== + +.. versionadded:: 1.12.0 + +.. caution:: unstable + +Kyuubi supports intercepting interactive statements on the server before they are routed to the engine, through custom statement interceptors. As a unified gateway, Kyuubi can host this governance once at the server layer instead of having every engine reimplement it, so administrators can inspect, reject, rewrite, or tune each statement without forking Kyuubi or writing the same logic for Spark, Flink, Trino, and JDBC separately. Typical use cases include rule-based SQL guards, calling an external policy service, risky-statement interception, auditing, SQL rewriting, and per-statement execution tuning. + +The interceptor sees statement text and gateway-level context, not an engine's analyzed plan or resolved objects. It can enforce text- or external-policy-based authorization, but it is not a replacement for semantic authorization such as Spark AuthZ. It also complements the existing extension points such as custom authentication, ``SessionConfAdvisor``, and event handlers. + +The Plugin Interface +-------------------- + +The public SPI lives in the ``kyuubi-server-plugin`` module under the ``org.apache.kyuubi.plugin`` package, alongside the existing ``SessionConfAdvisor`` and ``GroupProvider``. The module is dependency-free, so the interfaces only use JDK types and do not reference internal Kyuubi types, which keeps plugins decoupled from the server and the API stable. + +A single instance of each interceptor is created per Kyuubi server via its zero-arg constructor, and ``beforeExecuteStatement`` is invoked concurrently by many sessions. Implementations MUST be thread-safe and MUST NOT keep per-request mutable state in instance fields. + +.. code-block:: java + + public interface StatementInterceptor { + + // Called once when the server starts; conf is an immutable read-only snapshot + // of the full server configuration. Implementations read their own private keys. + default void initialize(Map conf) {} + + // Invoked for each statement before operation creation and engine routing. + StatementInterceptResult beforeExecuteStatement(StatementInterceptContext context); + + // Called at most once after initialization is attempted, either during failed startup + // or when the server stops. It must release resources allocated before initialize throws. + default void close() {} + } + +The context exposes a stable, gateway-level view of the statement using JDK types only. It intentionally does not expose the session configuration (which carries connection parameters and engine credentials) nor any mutable internal session object. + +.. code-block:: java + + public interface StatementInterceptContext { + String sessionId(); + String statementId(); // unique id, equal to the operation handle the client receives + String user(); // effective user the statement runs as (proxy user if impersonating) + String realUser(); // authenticated user before impersonation; equals user() if none + String ipAddress(); // empty string when unknown, never null + String statement(); // the current statement (after prior rewrites) + Map confOverlay(); // statement-level overlay, read-only + boolean runAsync(); + long queryTimeout(); // client-requested timeout in seconds, 0 means none + String engineType(); // upper-cased kyuubi.engine.type, e.g. SPARK_SQL / FLINK_SQL / TRINO + } + +An interceptor returns one of three decisions: + +.. code-block:: java + + public final class StatementInterceptResult { + public enum Action { PROCEED, REWRITE, REJECT } + + public static StatementInterceptResult proceed(); // keep the current statement + public static StatementInterceptResult proceed(Map conf); + public static StatementInterceptResult rewrite(String s); // replace it for the next interceptor and the engine + public static StatementInterceptResult rewrite(String s, Map conf); + public static StatementInterceptResult reject(String msg); // stop the chain and return an error to the client + public Map confOverlay(); // immutable config delta + } + +Enable Statement Interceptors +----------------------------- + +1. Create one or more classes implementing ``org.apache.kyuubi.plugin.StatementInterceptor``. +2. Compile and put the jar into ``$KYUUBI_HOME/jars``. +3. Add the configuration in ``kyuubi-defaults.conf``: + + .. code-block:: properties + + kyuubi.operation.statement.interceptors=com.example.SqlGuard,com.example.LlmSqlRewriter + +Interceptors run in the configured order. The execution semantics are: + +- the chain starts from the original statement; +- ``PROCEED`` keeps the current statement and passes it to the next interceptor; +- ``REWRITE`` replaces the current statement and passes the new one to the next interceptor and ultimately to the engine; +- ``PROCEED`` and ``REWRITE`` may also return per-statement config updates. Updates are accumulated in interceptor order, later values win, and each following interceptor sees the updated read-only overlay. The final overlay is passed to the engine operation and does not mutate the session configuration; +- ``REJECT`` stops the chain immediately, no operation is created, and the client receives an error carrying SQLState ``42501`` (insufficient privilege), so clients can tell a policy rejection from a generic syntax error; +- if an interceptor throws or returns ``null``, the statement fails (fail-closed). + +The interceptors are eagerly loaded and initialized at server startup, so a misconfigured or failing interceptor fails the server fast rather than at the first query. They are closed in reverse order when the server stops. + +Example +------- + +A simple guard that rejects statements containing configured keywords: + +.. code-block:: java + + package com.example; + + import java.util.*; + import java.util.stream.Collectors; + import org.apache.kyuubi.plugin.*; + + public class SqlGuard implements StatementInterceptor { + + private Set blockedKeywords; + + @Override + public void initialize(Map conf) { + String raw = conf.getOrDefault("example.sql.guard.blocked.keywords", "drop,truncate"); + blockedKeywords = Arrays.stream(raw.split(",")) + .map(s -> s.trim().toLowerCase(Locale.ROOT)) + .filter(s -> !s.isEmpty()) + .collect(Collectors.toSet()); + } + + @Override + public StatementInterceptResult beforeExecuteStatement(StatementInterceptContext ctx) { + String sql = ctx.statement().trim().toLowerCase(Locale.ROOT); + if (blockedKeywords.stream().anyMatch(sql::contains)) { + return StatementInterceptResult.reject("SQL is rejected by policy for user " + ctx.user()); + } + return StatementInterceptResult.proceed(); + } + } + +.. note:: The substring match above is for illustration only. A production SQL guard should use a parser, otherwise a statement like ``SELECT * FROM dropdown_events`` would be wrongly rejected by ``contains("drop")``. + +Then deploy the jar and enable it: + +.. code-block:: properties + + kyuubi.operation.statement.interceptors=com.example.SqlGuard + example.sql.guard.blocked.keywords=drop,truncate,delete + +Config Tuning Example +--------------------- + +Config tuning is orthogonal to rewriting: an interceptor may tune the current statement without changing its SQL, or do both in one result. For example, this interceptor reduces Spark shuffle parallelism for an interactive query: + +.. code-block:: java + + @Override + public StatementInterceptResult beforeExecuteStatement(StatementInterceptContext ctx) { + if (ctx.engineType().equals("SPARK_SQL") && isInteractive(ctx.statement())) { + return StatementInterceptResult.proceed( + Collections.singletonMap("spark.sql.shuffle.partitions", "32")); + } + return StatementInterceptResult.proceed(); + } + +Only configuration entries that the selected engine consumes at statement planning or execution time can take effect. Engine-launch or session-initialization settings are outside this SPI's per-statement scope. + +Notes +----- + +.. note:: The configuration is ``serverOnly``, so it cannot be overridden or disabled in a session; on the paths where interceptors run (see below), users cannot turn them off. + +.. note:: Interceptors only apply to the interactive statement path (``executeStatement`` on a SQL session), covering both engine-routed statements and server-side commands. Statements that never reach that path are not intercepted: batch jobs (submitted as applications), metadata operations such as ``getTables``, and the Data Agent REST endpoints (which forward the request text straight to the engine). + +.. caution:: ``beforeExecuteStatement`` runs synchronously on the statement-submission thread and adds directly to the submission latency. Interceptors that make external calls (for example to an authorization service or an LLM) must enforce their own timeout and retry limits, and choose their own degradation policy: governance interceptors should fail closed (``REJECT`` or throw), while enrichment interceptors may fail open (``PROCEED``). diff --git a/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptContext.java b/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptContext.java new file mode 100644 index 00000000000..24ee9fa432c --- /dev/null +++ b/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptContext.java @@ -0,0 +1,73 @@ +/* + * 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 java.util.Map; + +/** + * Stable, gateway-level context passed to a {@link StatementInterceptor}. All fields are JDK types. + * It intentionally does not expose the session configuration (which carries connection parameters + * and engine credentials) nor any mutable internal session object. + */ +public interface StatementInterceptContext { + + /** The session identifier the statement belongs to. */ + String sessionId(); + + /** + * The unique identifier of this statement, equal to the operation handle the client receives. It + * is allocated before interception and stays stable through the operation's whole lifecycle, so + * it can correlate the intercepted statement with its later operation and result set. + */ + String statementId(); + + /** + * The effective user the statement runs as. With impersonation enabled (for example {@code + * hive.server2.proxy.user}), this is the proxy user; otherwise it equals {@link #realUser()}. Use + * this as the identity for authorization and auditing. + */ + String user(); + + /** + * The real user that authenticated the connection, before any impersonation. Equals {@link + * #user()} when impersonation is not in effect. + */ + String realUser(); + + /** The client IP address; an empty string when unknown, never {@code null}. */ + String ipAddress(); + + /** The statement to be executed. With a chain of interceptors, this is the current statement. */ + String statement(); + + /** The per-statement configuration overlay (statement-level, read-only). */ + Map confOverlay(); + + /** Whether the statement is executed asynchronously. */ + boolean runAsync(); + + /** The client-requested query timeout in seconds; {@code 0} means no timeout. */ + long queryTimeout(); + + /** + * The engine type resolved from {@code kyuubi.engine.type}, upper-cased to the config's enum + * names such as {@code SPARK_SQL}, {@code FLINK_SQL}, {@code TRINO}, {@code HIVE_SQL}, {@code + * JDBC}. Match against these values, not lower-cased short names like {@code spark}. + */ + String engineType(); +} diff --git a/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptResult.java b/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptResult.java new file mode 100644 index 00000000000..ee8721495eb --- /dev/null +++ b/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptResult.java @@ -0,0 +1,128 @@ +/* + * 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 java.util.Collections; +import java.util.HashMap; +import java.util.Map; + +/** + * The decision returned by a {@link StatementInterceptor} for a single statement: proceed with the + * current statement, rewrite it for the next interceptor and the engine, or reject it. + */ +public final class StatementInterceptResult { + + public enum Action { + PROCEED, + REWRITE, + REJECT + } + + private final Action action; + private final String statement; // the rewritten statement when action is REWRITE + private final String message; // the rejection reason when action is REJECT + private final Map confOverlay; + + private StatementInterceptResult( + Action action, String statement, String message, Map confOverlay) { + this.action = action; + this.statement = statement; + this.message = message; + this.confOverlay = immutableConfOverlay(confOverlay); + } + + /** Keep the current statement and pass it to the next interceptor. */ + public static StatementInterceptResult proceed() { + return proceed(Collections.emptyMap()); + } + + /** + * Keep the current statement and apply {@code confOverlay} to this statement only. The entries + * are passed to subsequent interceptors and ultimately to the engine operation. + */ + public static StatementInterceptResult proceed(Map confOverlay) { + return new StatementInterceptResult(Action.PROCEED, null, null, confOverlay); + } + + /** + * Replace the current statement with {@code statement}. It is passed to the next interceptor and + * ultimately to the engine. {@code statement} must not be null or blank. + */ + public static StatementInterceptResult rewrite(String statement) { + return rewrite(statement, Collections.emptyMap()); + } + + /** + * Replace the current statement with {@code statement} and apply {@code confOverlay} to this + * statement only. Both are passed to subsequent interceptors and ultimately to the engine + * operation. {@code statement} must not be null or blank. + */ + public static StatementInterceptResult rewrite( + String statement, Map confOverlay) { + return new StatementInterceptResult( + Action.REWRITE, requireNonBlank(statement, "statement"), null, confOverlay); + } + + /** + * Reject the statement immediately; the chain stops and the client receives an error. {@code + * message} must not be null or blank. + */ + public static StatementInterceptResult reject(String message) { + return new StatementInterceptResult( + Action.REJECT, null, requireNonBlank(message, "message"), Collections.emptyMap()); + } + + public Action action() { + return action; + } + + public String statement() { + return statement; + } + + public String message() { + return message; + } + + /** The immutable per-statement configuration updates returned by this interceptor. */ + public Map confOverlay() { + return confOverlay; + } + + private static String requireNonBlank(String value, String name) { + if (value == null || value.trim().isEmpty()) { + throw new IllegalArgumentException(name + " must not be null or blank"); + } + return value; + } + + private static Map immutableConfOverlay(Map confOverlay) { + if (confOverlay == null) { + throw new IllegalArgumentException("confOverlay must not be null"); + } + Map copy = new HashMap<>(); + confOverlay.forEach( + (key, value) -> { + if (key == null || value == null) { + throw new IllegalArgumentException("confOverlay must not contain null keys or values"); + } + copy.put(key, value); + }); + return Collections.unmodifiableMap(copy); + } +} diff --git a/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptor.java b/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptor.java new file mode 100644 index 00000000000..5b98cf110f0 --- /dev/null +++ b/extensions/server/kyuubi-server-plugin/src/main/java/org/apache/kyuubi/plugin/StatementInterceptor.java @@ -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.plugin; + +import java.util.Map; + +/** + * A server-side extension point invoked before an interactive statement is routed to the engine. + * Implementations can inspect, reject, or rewrite each statement. + * + *

Each interceptor is instantiated once via its zero-arg constructor and held as a single global + * instance per Kyuubi server. {@link #beforeExecuteStatement} is called concurrently by many + * sessions/threads, so implementations MUST be thread-safe and MUST NOT keep per-request mutable + * state in instance fields. Resources prepared in {@link #initialize} should be shared read-only. + */ +public interface StatementInterceptor { + + /** + * Called once when the Kyuubi server starts. The passed map is an immutable read-only snapshot of + * the full server configuration taken at startup; it does not reflect later changes and must not + * be modified. Implementations read their own private keys from it. + */ + default void initialize(Map conf) {} + + /** Invoked for each statement before operation creation and engine routing. */ + StatementInterceptResult beforeExecuteStatement(StatementInterceptContext context); + + /** + * Called at most once after initialization is attempted, either when server startup fails or when + * the server stops. Implementations must also release resources allocated before {@link + * #initialize} throws. + */ + default void close() {} +} diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala index 7cb2dee3656..a78b131089c 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/ExecuteStatement.scala @@ -24,6 +24,7 @@ import scala.collection.JavaConverters._ import org.apache.hadoop.fs.Path import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.kyuubi.SparkDatasetHelper._ import org.apache.kyuubi.{KyuubiSQLException, Logging} @@ -37,6 +38,7 @@ import org.apache.kyuubi.session.Session class ExecuteStatement( session: Session, override val statement: String, + confOverlay: Map[String, String], override val shouldRunAsync: Boolean, queryTimeout: Long, incrementalCollect: Boolean, @@ -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) { + setState(OperationState.RUNNING) + info(diagnostics) + Thread.currentThread().setContextClassLoader(spark.sharedState.jarClassLoader) + addOperationListener() + result = spark.sql(statement) + iter = collectAsIterator(result) + setCompiledStateIfNeeded() + setState(OperationState.FINISHED) + } } } catch { onError(cancel = true) @@ -212,6 +220,7 @@ class ExecuteStatement( class ArrowBasedExecuteStatement( session: Session, override val statement: String, + confOverlay: Map[String, String], override val shouldRunAsync: Boolean, queryTimeout: Long, incrementalCollect: Boolean, @@ -219,6 +228,7 @@ class ArrowBasedExecuteStatement( extends ExecuteStatement( session, statement, + confOverlay, shouldRunAsync, queryTimeout, incrementalCollect, diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala index b48863e3c8f..4a6e0520849 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/PlanOnlyStatement.scala @@ -46,6 +46,7 @@ import org.apache.kyuubi.util.reflect.DynMethods class PlanOnlyStatement( session: Session, override val statement: String, + confOverlay: Map[String, String], mode: PlanOnlyMode, override protected val handle: OperationHandle) extends SparkOperation(session) { @@ -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) parsed match { case cmd if planExcludes.contains(cmd.getClass.getSimpleName) => diff --git a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala index 5533d9c45e2..706e9b398b6 100644 --- a/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala +++ b/externals/kyuubi-spark-sql-engine/src/main/scala/org/apache/kyuubi/engine/spark/operation/SparkSQLOperationManager.scala @@ -95,6 +95,7 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n new ArrowBasedExecuteStatement( session, statement, + confOverlay, runAsync, queryTimeout, incrementalCollect, @@ -103,13 +104,14 @@ class SparkSQLOperationManager private (name: String) extends OperationManager(n new ExecuteStatement( session, statement, + confOverlay, runAsync, queryTimeout, incrementalCollect, opHandle) } case mode => - new PlanOnlyStatement(session, statement, mode, opHandle) + new PlanOnlyStatement(session, statement, confOverlay, mode, opHandle) } case OperationLanguages.SCALA => val repl = sessionToRepl.getOrElseUpdate(session.handle, KyuubiSparkILoop(spark)) diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala index 4239f8b6187..a21f77af375 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala @@ -3195,6 +3195,20 @@ object KyuubiConf { } .createWithDefault("hadoop") + val STATEMENT_INTERCEPTORS: OptionalConfigEntry[Seq[String]] = + buildConf("kyuubi.operation.statement.interceptors") + .doc("A comma-separated list of statement interceptor plugins for Kyuubi Server. Each " + + "value should be a subclass of `org.apache.kyuubi.plugin.StatementInterceptor` with a " + + "zero-arg constructor. They are invoked in the configured order on the server before " + + "each interactive statement is routed to the engine, and can inspect, reject, or " + + "rewrite the statement.") + .version("1.12.0") + .audience(SERVER) + .immutable + .stringConf + .toSequence() + .createOptional + val SERVER_NAME: OptionalConfigEntry[String] = buildConf("kyuubi.server.name") .doc("The name of Kyuubi Server.") diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala index 02c22e4f6a7..98ed07e7918 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecuteStatement.scala @@ -36,8 +36,9 @@ class ExecuteStatement( override val statement: String, confOverlay: Map[String, String], override val shouldRunAsync: Boolean, - queryTimeout: Long) - extends KyuubiOperation(session) { + queryTimeout: Long, + operationHandle: OperationHandle = OperationHandle()) + extends KyuubiOperation(session, operationHandle) { final private val _operationLog: OperationLog = if (shouldRunAsync) { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala index a59c2db7b77..e7c02266d51 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/ExecutedCommandExec.scala @@ -27,8 +27,9 @@ import org.apache.kyuubi.sql.schema.SchemaHelper class ExecutedCommandExec( session: KyuubiSessionImpl, override val shouldRunAsync: Boolean, - command: RunnableCommand) - extends KyuubiOperation(session) { + command: RunnableCommand, + operationHandle: OperationHandle = OperationHandle()) + extends KyuubiOperation(session, operationHandle) { private lazy val _operationLog: OperationLog = if (shouldRunAsync) { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala index 897c90c2396..9bffbc84acd 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperation.scala @@ -35,7 +35,15 @@ import org.apache.kyuubi.shaded.thrift.TException import org.apache.kyuubi.shaded.thrift.transport.TTransportException import org.apache.kyuubi.util.ThriftUtils -abstract class KyuubiOperation(session: Session) extends AbstractOperation(session) { +// The handle is injected as a constructor paramaccessor `override protected val handle` (mirroring +// the Spark engine operations) so the server can pass a pre-allocated handle whose id the statement +// interceptor already observed. It must stay a constructor parameter, not a body `val`, so that +// AbstractOperation's `statementId = handle.identifier.toString` reads the injected handle during +// superclass initialization rather than a null field. +abstract class KyuubiOperation( + session: Session, + override protected val handle: OperationHandle = OperationHandle()) + extends AbstractOperation(session) { MetricsSystem.tracing { ms => ms.incCount(MetricRegistry.name(OPERATION_OPEN, opType)) diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala index a248fe2a832..197627f1612 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/KyuubiOperationManager.scala @@ -58,17 +58,39 @@ class KyuubiOperationManager private (name: String) extends OperationManager(nam statement: String, confOverlay: Map[String, String], runAsync: Boolean, - queryTimeout: Long): Operation = { + queryTimeout: Long): Operation = + newExecuteStatementOperation( + session, + statement, + confOverlay, + runAsync, + queryTimeout, + OperationHandle()) + + def newExecuteStatementOperation( + session: Session, + statement: String, + confOverlay: Map[String, String], + runAsync: Boolean, + queryTimeout: Long, + operationHandle: OperationHandle): Operation = { val operation = - new ExecuteStatement(session, statement, confOverlay, runAsync, getQueryTimeout(queryTimeout)) + new ExecuteStatement( + session, + statement, + confOverlay, + runAsync, + getQueryTimeout(queryTimeout), + operationHandle) addOperation(operation) } def newExecuteOnServerOperation( session: KyuubiSessionImpl, runAsync: Boolean, - command: RunnableCommand): Operation = { - val operation = new ExecutedCommandExec(session, runAsync, command) + command: RunnableCommand, + operationHandle: OperationHandle = OperationHandle()): Operation = { + val operation = new ExecutedCommandExec(session, runAsync, command, operationHandle) addOperation(operation) } diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/PluginLoader.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/PluginLoader.scala index 1bc80dc7da1..ce7428b15c3 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/PluginLoader.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/PluginLoader.scala @@ -44,6 +44,24 @@ private[kyuubi] object PluginLoader { } } + def loadStatementInterceptors(conf: KyuubiConf): Seq[StatementInterceptor] = { + conf.get(KyuubiConf.STATEMENT_INTERCEPTORS).getOrElse(Nil).map { interceptorClassName => + try { + DynConstructors.builder.impl(interceptorClassName) + .buildChecked[StatementInterceptor].newInstance() + } catch { + case _: ClassCastException => + throw new KyuubiException( + s"Class $interceptorClassName is not a child of " + + s"'${classOf[StatementInterceptor].getName}'.") + case NonFatal(e) => + throw new IllegalArgumentException( + s"Error while instantiating '$interceptorClassName': ", + e) + } + } + } + def loadGroupProvider(conf: KyuubiConf): GroupProvider = { val groupProviderClass = conf.get(KyuubiConf.GROUP_PROVIDER) try { diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/StatementInterception.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/plugin/StatementInterception.scala new file mode 100644 index 00000000000..89dd8983b6f --- /dev/null +++ b/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 + } + InterceptedStatement(currentStatement, currentConfOverlay) + } +} diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala index 436b610861b..3a16e1032f3 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala +++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionImpl.scala @@ -313,15 +313,38 @@ class KyuubiSessionImpl( confOverlay: Map[String, String], runAsync: Boolean, queryTimeout: Long): OperationHandle = withAcquireRelease() { - val kyuubiNode = parser.parsePlan(statement) + // Allocate the operation handle up front so the interceptor sees the same id the operation + // and its result set will carry, then reuse it instead of letting the operation self-generate. + val operationHandle = OperationHandle() + val interception = sessionManager.interceptStatement( + handle.identifier.toString, + operationHandle.identifier.toString, + user, + realUser, + ipAddress, + statement, + confOverlay, + runAsync, + queryTimeout, + sessionConf.get(ENGINE_TYPE)) + val kyuubiNode = parser.parsePlan(interception.statement) kyuubiNode match { case command: RunnableCommand => val operation = sessionManager.operationManager.newExecuteOnServerOperation( this, runAsync, - command) + command, + operationHandle) + runOperation(operation) + case _ => + val operation = sessionManager.operationManager.newExecuteStatementOperation( + this, + interception.statement, + interception.confOverlay, + runAsync, + queryTimeout, + operationHandle) runOperation(operation) - case _ => super.executeStatement(statement, confOverlay, runAsync, queryTimeout) } } 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..6ca7404ad41 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 @@ -20,6 +20,7 @@ package org.apache.kyuubi.session import java.util.concurrent.{Semaphore, TimeUnit} import scala.collection.JavaConverters._ +import scala.util.control.NonFatal import com.codahale.metrics.MetricRegistry import com.google.common.annotations.VisibleForTesting @@ -35,7 +36,7 @@ import org.apache.kyuubi.engine.KyuubiApplicationManager import org.apache.kyuubi.metrics.MetricsConstants._ import org.apache.kyuubi.metrics.MetricsSystem import org.apache.kyuubi.operation.{KyuubiOperationManager, OperationState} -import org.apache.kyuubi.plugin.{GroupProvider, PluginLoader, SessionConfAdvisor} +import org.apache.kyuubi.plugin.{GroupProvider, InterceptedStatement, PluginLoader, SessionConfAdvisor, StatementInterceptContextImpl, StatementInterception, StatementInterceptor} import org.apache.kyuubi.server.metadata.{MetadataManager, MetadataRequestsRetryRef} import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter} import org.apache.kyuubi.service.TempFileService @@ -62,6 +63,11 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { lazy val sessionConfAdvisor: Seq[SessionConfAdvisor] = PluginLoader.loadSessionConfAdvisor(conf) lazy val groupProvider: GroupProvider = PluginLoader.loadGroupProvider(conf) + // Statement interceptors are eagerly loaded and initialized in initialize(conf), not lazy like + // the plugins above, so that a misconfigured interceptor fails the server at startup. Volatile + // because it is written once in initialize(conf) and read from many frontend request threads. + @volatile private var statementInterceptors: Seq[StatementInterceptor] = Nil + private var limiter: Option[SessionLimiter] = None private var batchLimiter: Option[SessionLimiter] = None lazy val (signingPrivateKey, signingPublicKey) = SignUtils.generateKeyPair() @@ -84,6 +90,79 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) { initSessionLimiter(conf) initEngineStartupProcessSemaphore(conf) super.initialize(conf) + // Initialize interceptors after super.initialize so a failure there does not leave them open. + // StatementInterception.initialize closes every attempted interceptor on its own + // failure, and the field is assigned only after all succeed, so stop() never double-closes. + val loadedInterceptors = PluginLoader.loadStatementInterceptors(conf) + // conf.getAll returns a fresh immutable snapshot whose asJava view is already read-only. + val serverConf: java.util.Map[String, String] = conf.getAll.asJava + StatementInterception.initialize(loadedInterceptors, serverConf) + statementInterceptors = loadedInterceptors + } + + override def stop(): Unit = { + try { + super.stop() + } finally { + // Release interceptor external resources after services/sessions are stopped. + closeStatementInterceptors() + } + } + + private def closeStatementInterceptors(): Unit = { + val interceptorsToClose = synchronized { + val current = statementInterceptors + statementInterceptors = Nil + current + } + interceptorsToClose.reverse.foreach { interceptor => + try { + interceptor.close() + } catch { + case NonFatal(e) => + warn(s"Error closing statement interceptor ${interceptor.getClass.getName}", e) + } + } + } + + /** + * Run the configured statement interceptors before an interactive statement is routed to the + * engine. Returns the (possibly rewritten) statement and per-statement config overlay, or throws + * if an interceptor rejects it, throws, or returns null. + */ + def interceptStatement( + sessionId: String, + statementId: String, + user: String, + realUser: String, + ipAddress: String, + statement: String, + confOverlay: Map[String, String], + runAsync: Boolean, + queryTimeout: Long, + engineType: String): InterceptedStatement = { + if (statementInterceptors.isEmpty) { + InterceptedStatement(statement, confOverlay) + } else { + val safeIpAddress = Option(ipAddress).getOrElse("") + val safeEngineType = Option(engineType).getOrElse("") + StatementInterception.run( + statementInterceptors, + statement, + confOverlay, + (currentStatement, currentConfOverlay) => + new StatementInterceptContextImpl( + sessionId, + statementId, + user, + realUser, + safeIpAddress, + currentStatement, + currentConfOverlay, + runAsync, + queryTimeout, + safeEngineType)) + } } override protected def createSession( diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/StatementInterceptorSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/StatementInterceptorSuite.scala new file mode 100644 index 00000000000..ab34f179a19 --- /dev/null +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/operation/StatementInterceptorSuite.scala @@ -0,0 +1,188 @@ +/* + * 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.operation + +import java.sql.SQLException +import java.util.Collections +import java.util.Locale + +import org.apache.kyuubi.{Utils, WithKyuubiServer} +import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.plugin.{StatementInterceptContext, StatementInterceptor, StatementInterceptResult} +import org.apache.kyuubi.shaded.hive.service.rpc.thrift.{TExecuteStatementReq, TFetchOrientation, TFetchResultsReq, TStatusCode} + +/** + * End-to-end coverage of statement interceptors over the real JDBC -> server -> engine path. + * Requires a packaged engine, so it runs in the engine-backed CI jobs rather than a bare unit run. + */ +class StatementInterceptorSuite extends WithKyuubiServer with HiveJDBCTestHelper { + + override protected def jdbcUrl: String = + s"jdbc:kyuubi://${server.frontendServices.head.connectionUrl}/;" + override protected val URL_PREFIX: String = "jdbc:kyuubi://" + + override protected val conf: KyuubiConf = { + KyuubiConf() + .set(KyuubiConf.ENGINE_SHARE_LEVEL, "connection") + .set(KyuubiConf.ENGINE_SPARK_MAX_INITIAL_WAIT.key, "0") + // allow the test user to impersonate, so the effective user can differ from the real user + .set(s"hadoop.proxyuser.${Utils.currentUser}.groups", "*") + .set(s"hadoop.proxyuser.${Utils.currentUser}.hosts", "*") + .set( + KyuubiConf.STATEMENT_INTERCEPTORS, + Seq(classOf[TestStatementInterceptor].getName)) + } + + test("PROCEED - statement executes unchanged") { + withJdbcStatement() { statement => + val rs = statement.executeQuery("SELECT 1 AS id") + assert(rs.next()) + assert(rs.getInt("id") === 1) + assert(!rs.next()) + } + } + + test("REWRITE - the rewritten statement is executed") { + withJdbcStatement() { statement => + val rs = statement.executeQuery("SELECT 'rewrite_me' AS result") + assert(rs.next()) + assert(rs.getString("result") === "rewritten") + assert(!rs.next()) + } + } + + test("TUNE - interceptor config overrides apply to this statement") { + withJdbcStatement() { statement => + statement.execute("SET spark.sql.shuffle.partitions=11") + + val rs = statement.executeQuery("SELECT 'tune_conf' AS marker") + val tunedPlan = + Iterator.continually(rs).takeWhile(_.next()).map(_.getString(1)).mkString("\n") + assert(tunedPlan.contains("hashpartitioning")) + assert(tunedPlan.contains(", 7)")) + + val nextRs = statement.executeQuery( + "EXPLAIN FORMATTED SELECT * FROM range(100) DISTRIBUTE BY id") + val nextPlan = + Iterator.continually(nextRs).takeWhile(_.next()).map(_.getString(1)).mkString("\n") + assert(nextPlan.contains("hashpartitioning")) + assert(nextPlan.contains(", 11)")) + } + } + + test("interceptor config overlay applies in plan-only mode without leaking") { + withJdbcStatement() { statement => + statement.execute("SET spark.sql.shuffle.partitions=11") + statement.execute(s"SET ${KyuubiConf.OPERATION_PLAN_ONLY_MODE.key}=${PhysicalMode.name}") + + val rs = statement.executeQuery("SELECT 'plan_only_conf' AS marker") + assert(rs.next()) + val interceptedPlan = rs.getString(1).toLowerCase(Locale.ROOT) + assert(interceptedPlan.contains("hashpartitioning")) + assert(interceptedPlan.contains(", 7)")) + + val nextRs = statement.executeQuery("SELECT * FROM range(100) DISTRIBUTE BY id") + assert(nextRs.next()) + val nextPlan = nextRs.getString(1).toLowerCase(Locale.ROOT) + assert(nextPlan.contains("hashpartitioning")) + assert(nextPlan.contains(", 11)")) + } + } + + test("REJECT - an engine-routed statement returns an error to the client") { + withJdbcStatement() { statement => + val e = intercept[SQLException] { + statement.executeQuery("SELECT 'forbidden' AS col") + } + assert(e.getMessage.contains("blocked: forbidden keyword")) + } + } + + test("REJECT - a server-side command (executeOnServer path) is also intercepted") { + withJdbcStatement() { statement => + val e = intercept[SQLException] { + statement.execute("KYUUBI DESC SESSION") + } + assert(e.getMessage.contains("blocked: server-side command")) + } + } + + test("the statement id is exposed to the interceptor") { + withSessionHandle { (client, sessionHandle) => + val request = new TExecuteStatementReq(sessionHandle, "SELECT 'echo_stmt_id' AS x") + request.setRunAsync(false) + val response = client.ExecuteStatement(request) + assert(response.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS) + + val operationHandle = response.getOperationHandle + val fetchRequest = new TFetchResultsReq( + operationHandle, + TFetchOrientation.FETCH_NEXT, + 1) + val fetchResponse = client.FetchResults(fetchRequest) + assert(fetchResponse.getStatus.getStatusCode === TStatusCode.SUCCESS_STATUS) + val statementId = fetchResponse.getResults.getColumns.get(0).getStringVal.getValues.get(0) + assert(statementId === OperationHandle(operationHandle).identifier.toString) + } + } + + test("the interceptor sees the effective user and the real user under impersonation") { + val proxyUser = "proxy_alice" + assert(proxyUser !== Utils.currentUser) + withSessionConf(Map("hive.server2.proxy.user" -> proxyUser))(Map.empty)(Map.empty) { + withJdbcStatement() { statement => + val rs = statement.executeQuery("SELECT 'echo_users' AS x") + assert(rs.next()) + // The interceptor rewrote the query to echo context.user() and context.realUser(). + assert(rs.getString("effective_user") === proxyUser) + assert(rs.getString("real_user") === Utils.currentUser) + assert(rs.getString("effective_user") !== rs.getString("real_user")) + assert(!rs.next()) + } + } + } +} + +class TestStatementInterceptor extends StatementInterceptor { + override def beforeExecuteStatement( + context: StatementInterceptContext): StatementInterceptResult = { + val sql = context.statement().trim.toLowerCase(Locale.ROOT) + if (sql.startsWith("kyuubi ")) { + StatementInterceptResult.reject("blocked: server-side command") + } else if (sql.contains("forbidden")) { + StatementInterceptResult.reject("blocked: forbidden keyword") + } else if (sql.contains("rewrite_me")) { + StatementInterceptResult.rewrite("SELECT 'rewritten' AS result") + } else if (sql.contains("tune_conf")) { + StatementInterceptResult.rewrite( + "EXPLAIN FORMATTED SELECT * FROM range(100) DISTRIBUTE BY id", + Collections.singletonMap("spark.sql.shuffle.partitions", "7")) + } else if (sql.contains("plan_only_conf")) { + StatementInterceptResult.rewrite( + "SELECT * FROM range(100) DISTRIBUTE BY id", + Collections.singletonMap("spark.sql.shuffle.partitions", "7")) + } else if (sql.contains("echo_stmt_id")) { + StatementInterceptResult.rewrite(s"SELECT '${context.statementId()}' AS sid") + } else if (sql.contains("echo_users")) { + StatementInterceptResult.rewrite( + s"SELECT '${context.user()}' AS effective_user, '${context.realUser()}' AS real_user") + } else { + StatementInterceptResult.proceed() + } + } +} diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/PluginLoaderSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/PluginLoaderSuite.scala index fa4505cc248..9047e20d9cb 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/PluginLoaderSuite.scala +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/PluginLoaderSuite.scala @@ -120,6 +120,28 @@ class PluginLoaderSuite extends KyuubiFunSuite { assert(msg2.startsWith("Error while instantiating 'non.exists'")) } + test("StatementInterceptor - load, order and wrong class") { + val conf = new KyuubiConf(false) + assert(PluginLoader.loadStatementInterceptors(conf).isEmpty) + + conf.set( + KyuubiConf.STATEMENT_INTERCEPTORS, + Seq(classOf[NoopStatementInterceptor].getName, classOf[NoopStatementInterceptor].getName)) + assert(PluginLoader.loadStatementInterceptors(conf).size === 2) + + conf.set(KyuubiConf.STATEMENT_INTERCEPTORS, Seq(classOf[InvalidStatementInterceptor].getName)) + val msg1 = intercept[KyuubiException] { + PluginLoader.loadStatementInterceptors(conf) + }.getMessage + assert(msg1.contains(s"is not a child of '${classOf[StatementInterceptor].getName}'")) + + conf.set(KyuubiConf.STATEMENT_INTERCEPTORS, Seq("non.exists")) + val msg2 = intercept[IllegalArgumentException] { + PluginLoader.loadStatementInterceptors(conf) + }.getMessage + assert(msg2.startsWith("Error while instantiating 'non.exists'")) + } + test("HadoopGroupProvider") { val conf = new KyuubiConf(false) conf.set(KyuubiConf.GROUP_PROVIDER, "hadoop") @@ -133,6 +155,13 @@ class PluginLoaderSuite extends KyuubiFunSuite { class InvalidSessionConfAdvisor class InvalidGroupProvider +class InvalidStatementInterceptor + +class NoopStatementInterceptor extends StatementInterceptor { + override def beforeExecuteStatement( + context: StatementInterceptContext): StatementInterceptResult = + StatementInterceptResult.proceed() +} class TestSessionConfAdvisor extends SessionConfAdvisor { override def getConfOverlay( diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/StatementInterceptionSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/StatementInterceptionSuite.scala new file mode 100644 index 00000000000..cabf8a3910b --- /dev/null +++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/plugin/StatementInterceptionSuite.scala @@ -0,0 +1,235 @@ +/* + * 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.concurrent.{Await, Future} +import scala.concurrent.ExecutionContext.Implicits.global +import scala.concurrent.duration._ + +import org.apache.kyuubi.{KyuubiFunSuite, KyuubiSQLException} + +class StatementInterceptionSuite extends KyuubiFunSuite { + + private def contextFor( + statement: String, + confOverlay: java.util.Map[String, String]): StatementInterceptContext = + new StatementInterceptContextImpl( + "session-1", + "op-1", + "alice", + "alice_real", + "127.0.0.1", + statement, + confOverlay, + true, + 0L, + "SPARK_SQL") + + private def runResult( + interceptors: Seq[StatementInterceptor], + statement: String, + confOverlay: Map[String, String] = Map.empty): InterceptedStatement = + StatementInterception.run(interceptors, statement, confOverlay, contextFor) + + private def run(interceptors: Seq[StatementInterceptor], statement: String): String = + runResult(interceptors, statement).statement + + private def interceptor( + f: StatementInterceptContext => StatementInterceptResult): StatementInterceptor = + new StatementInterceptor { + override def beforeExecuteStatement( + context: StatementInterceptContext): StatementInterceptResult = f(context) + } + + test("no interceptors returns the original statement") { + assert(run(Nil, "SELECT 1") === "SELECT 1") + } + + test("PROCEED keeps the statement unchanged") { + val i = interceptor(_ => StatementInterceptResult.proceed()) + assert(run(Seq(i), "SELECT 1") === "SELECT 1") + } + + test("REWRITE replaces the statement") { + val i = interceptor(_ => StatementInterceptResult.rewrite("SELECT 2")) + assert(run(Seq(i), "SELECT 1") === "SELECT 2") + } + + test("PROCEED can tune the per-statement config overlay") { + val tuning = Map("spark.sql.shuffle.partitions" -> "32").asJava + val i = interceptor(_ => StatementInterceptResult.proceed(tuning)) + val result = runResult(Seq(i), "SELECT 1") + assert(result.statement === "SELECT 1") + assert(result.confOverlay === tuning.asScala) + } + + test("REWRITE can also tune the per-statement config overlay") { + val i = interceptor(_ => + StatementInterceptResult.rewrite( + "SELECT 2", + Map("spark.sql.adaptive.enabled" -> "true").asJava)) + val result = runResult(Seq(i), "SELECT 1") + assert(result.statement === "SELECT 2") + assert(result.confOverlay === Map("spark.sql.adaptive.enabled" -> "true")) + } + + test("config tuning chains in order and later values win") { + var seenBySecond: Map[String, String] = Map.empty + val first = interceptor(_ => + StatementInterceptResult.proceed( + Map("shared" -> "first", "first-only" -> "1").asJava)) + val second = interceptor { context => + seenBySecond = context.confOverlay().asScala.toMap + StatementInterceptResult.proceed( + Map("shared" -> "second", "second-only" -> "2").asJava) + } + val result = runResult( + Seq(first, second), + "SELECT 1", + Map("shared" -> "client", "client-only" -> "0")) + assert(seenBySecond === Map( + "shared" -> "first", + "client-only" -> "0", + "first-only" -> "1")) + assert(result.confOverlay === Map( + "shared" -> "second", + "client-only" -> "0", + "first-only" -> "1", + "second-only" -> "2")) + } + + test("result config overlay is a defensive immutable copy") { + val source = new java.util.HashMap[String, String]() + source.put("key", "value") + val result = StatementInterceptResult.proceed(source) + source.put("key", "changed") + assert(result.confOverlay().get("key") === "value") + intercept[UnsupportedOperationException] { + result.confOverlay().put("other", "value") + } + } + + test("REJECT throws KyuubiSQLException carrying the message and a policy SQLState") { + val i = interceptor(_ => StatementInterceptResult.reject("blocked by policy")) + val e = intercept[KyuubiSQLException](run(Seq(i), "SELECT 1")) + assert(e.getMessage.contains("blocked by policy")) + assert(e.getSQLState === "42501") + } + + test("interceptors run in order and rewrites chain") { + val appendA = interceptor(ctx => StatementInterceptResult.rewrite(ctx.statement() + " a")) + val appendB = interceptor(ctx => StatementInterceptResult.rewrite(ctx.statement() + " b")) + assert(run(Seq(appendA, appendB), "x") === "x a b") + } + + test("later interceptors see the rewritten statement") { + var seenBySecond: String = null + val first = interceptor(_ => StatementInterceptResult.rewrite("REWRITTEN")) + val second = interceptor { ctx => + seenBySecond = ctx.statement() + StatementInterceptResult.proceed() + } + run(Seq(first, second), "ORIGINAL") + assert(seenBySecond === "REWRITTEN") + } + + test("the context exposes the statement id") { + var seenStatementId: String = null + val i = interceptor { ctx => + seenStatementId = ctx.statementId() + StatementInterceptResult.proceed() + } + run(Seq(i), "SELECT 1") + assert(seenStatementId === "op-1") + } + + test("the context exposes the effective and real user") { + var seenUser: String = null + var seenRealUser: String = null + val i = interceptor { ctx => + seenUser = ctx.user() + seenRealUser = ctx.realUser() + StatementInterceptResult.proceed() + } + run(Seq(i), "SELECT 1") + assert(seenUser === "alice") + assert(seenRealUser === "alice_real") + } + + test("a throwing interceptor fails the statement") { + val i = interceptor(_ => throw new RuntimeException("boom")) + val e = intercept[KyuubiSQLException](run(Seq(i), "SELECT 1")) + assert(e.getMessage.contains("failed")) + } + + test("a null result fails the statement") { + val i = interceptor(_ => null) + val e = intercept[KyuubiSQLException](run(Seq(i), "SELECT 1")) + assert(e.getMessage.contains("returned null")) + } + + test("REJECT short-circuits later interceptors") { + val rejecting = interceptor(_ => StatementInterceptResult.reject("blocked")) + val mustNotRun = interceptor(_ => throw new RuntimeException("must not run")) + val e = intercept[KyuubiSQLException](run(Seq(rejecting, mustNotRun), "SELECT 1")) + assert(e.getMessage.contains("blocked")) + } + + test("rewrite/reject reject null or blank arguments") { + intercept[IllegalArgumentException](StatementInterceptResult.rewrite(null)) + intercept[IllegalArgumentException](StatementInterceptResult.rewrite(" ")) + intercept[IllegalArgumentException](StatementInterceptResult.reject(null)) + intercept[IllegalArgumentException](StatementInterceptResult.reject("")) + intercept[IllegalArgumentException](StatementInterceptResult.proceed(null)) + intercept[IllegalArgumentException](StatementInterceptResult.rewrite("SELECT 1", null)) + } + + test("initialize closes attempted interceptors in reverse order on failure") { + val closed = new java.util.concurrent.CopyOnWriteArrayList[String]() + def okInterceptor(name: String): StatementInterceptor = new StatementInterceptor { + override def beforeExecuteStatement( + context: StatementInterceptContext): StatementInterceptResult = + StatementInterceptResult.proceed() + override def close(): Unit = closed.add(name) + } + val failing = new StatementInterceptor { + override def initialize(conf: java.util.Map[String, String]): Unit = + throw new RuntimeException("init failed") + override def beforeExecuteStatement( + context: StatementInterceptContext): StatementInterceptResult = + StatementInterceptResult.proceed() + override def close(): Unit = closed.add("failing") + } + val e = intercept[RuntimeException] { + StatementInterception.initialize( + Seq(okInterceptor("a"), okInterceptor("b"), failing), + java.util.Collections.emptyMap[String, String]()) + } + assert(e.getMessage.contains("init failed")) + assert(closed.asScala.toList === List("failing", "b", "a")) + } + + test("run is thread-safe under concurrent invocation") { + val appendX = interceptor(ctx => StatementInterceptResult.rewrite(ctx.statement() + "-x")) + val futures = (1 to 200).map(i => Future(run(Seq(appendX), s"q$i"))) + val results = Await.result(Future.sequence(futures), 30.seconds) + assert(results.forall(_.endsWith("-x"))) + assert(results.distinct.size === 200) + } +}