[KYUUBI #7458] Add server-side statement interceptor SPI

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 |

docs/extensions/server/index.rst

@@ -25,5 +25,6 @@ ability of kyuubi servers.
 
     authentication
     configuration
+    statement_interceptor
     events
     applications

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<String, String> 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<String, String> 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<String, String> conf);
+     public static StatementInterceptResult rewrite(String s);  // replace it for the next interceptor and the engine
+     public static StatementInterceptResult rewrite(String s, Map<String, String> conf);
+     public static StatementInterceptResult reject(String msg); // stop the chain and return an error to the client
+     public Map<String, String> 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<String> blockedKeywords;
+
+     @Override
+     public void initialize(Map<String, String> 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``).

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<String, String> 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();
+}