根据条件动态修改SQL，添加指定过滤条件

最近生产项目遇到一个奇葩请求，要针对某一类型数据进行过滤，又不能影响原业务数据，所以考虑通过Mybatis插件实现动态改写，记录改写内容，方便记忆和学习。

一共创建了4个Java类，及部分配置：

DataScopeFilter：

package com.common.filter;

import org.apache.commons.lang3.StringUtils;

import javax.servlet.Filter;
import javax.servlet.FilterChain;
import javax.servlet.FilterConfig;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import java.io.IOException;

/**
 * 读取请求头中的数据范围过滤开关，存入 ThreadLocal 供 MyBatis 插件使用。
 *
 * <p>本 Filter 不是 Spring Bean（由 RootConfigurer 手工创建后塞进 FilterRegistrationBean），
 * 因此请求头名通过构造参数传入，不使用 @Value。</p>
 */
public class DataScopeFilter implements Filter {

    /** 开关请求头名 */
    private final String headerName;

    public DataScopeFilter(String headerName) {
        this.headerName = headerName;
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        try {
            if (request instanceof HttpServletRequest) {
                String value = ((HttpServletRequest) request).getHeader(headerName);
                if (StringUtils.isNotBlank(value)) {
                    DataScopeHolder.set(value);
                }
            }
            chain.doFilter(request, response);
        } finally {
            // 容器线程池会复用线程，不清理会让开关值泄漏到下一个请求
            DataScopeHolder.clear();
        }
    }

    @Override
    public void destroy() {
    }
}

 

DataScopeHolder：

package com.common.filter;

/**
 * 数据范围过滤开关值的 ThreadLocal 载体。
 *
 * <p>必须提供 clear()：容器线程池会复用线程，不清理会让上一次请求的开关值
 * 泄漏到下一次请求。capital-core 的 MultiTenantContextHolder 缺少该方法，不要照抄。</p>
 */
public final class DataScopeHolder {

    private static final ThreadLocal<String> HOLDER = new ThreadLocal<String>();

    private DataScopeHolder() {
    }

    public static void set(String value) {
        HOLDER.set(value);
    }

    public static String get() {
        return HOLDER.get();
    }

    public static void clear() {
        HOLDER.remove();
    }
}

 

DataScopeInterceptor：

package com.common.filter;

import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.scripting.xmltags.DynamicContext;
import org.apache.ibatis.scripting.xmltags.DynamicSqlSource;
import org.apache.ibatis.scripting.xmltags.SqlNode;
import org.apache.ibatis.session.Configuration;
import java.util.Map;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.HashSet;
import java.util.Locale;
import java.util.Properties;
import java.util.Set;

/**
 * 数据范围过滤条件的 MyBatis 插件。
 *
 * <p>拦截 Executor.query，把配置的条件追加到白名单表的查询上。</p>
 *
 * <p>插件必须在 mybatis-config.xml 的 &lt;plugins&gt; 中声明在 PagerInterceptor 之前（更内层）：
 * MyBatis 的插件顺序是反的，声明在最后的先执行。PagerInterceptor 内部执行分页 COUNT 查询时
 * 走 invocation.proceed()，目标是 Executor；若本插件在最外层，COUNT 查询会绕过它，
 * 导致「列表被过滤、count 没被过滤」，列表与总数口径不一致
 * （mapper 里预定义了 xxx_COUNT 语句时尤其明显）。</p>
 */
@Intercepts(@Signature(type = Executor.class, method = "query",
        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}))
public class DataScopeInterceptor implements Interceptor {

    private static final Logger logger = LoggerFactory.getLogger(DataScopeInterceptor.class);

    /** 白名单表名（小写） */
    private Set<String> tables = Collections.emptySet();

    /** 过滤列名 */
    private String column;

    /** 条件模板，含 {col} 占位符 */
    private String condition;

    /** SQL 方言 */
    private String dialect = "mariadb";

    /** 插件是否可用；配置非法时置为 false，全部查询原样放行 */
    private boolean enabled = true;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        if (!enabled) {
            return invocation.proceed();
        }
        String switchValue = DataScopeHolder.get();
        if (StringUtils.isBlank(switchValue)) {
            return invocation.proceed();
        }
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        SqlSource newSqlSource = rewritingSqlSource(ms, args[1]);
        if (newSqlSource == null) {
            return invocation.proceed();
        }
        args[0] = newMappedStatement(ms, newSqlSource);
        return invocation.proceed();
    }

    /**
     * 构造一个会在求值时改写 SQL 的 SqlSource。
     *
     * <p>动态 mapper 必须保持 DynamicSqlSource：capital-core 的 PagerInterceptor 靠
     * {@code sqlSource instanceof DynamicSqlSource} 判断走动态分支，动态分支会用
     * PagerDynamicSqlSource 重新求值并把 &lt;foreach&gt; 展开值复制回 BoundSql。
     * 若换成非动态的 SqlSource，PagerInterceptor 会走 Static 分支重建 BoundSql，
     * 展开值就丢了，mapper 会在取参数时报 "no getter for property '__frch_item_0'"。</p>
     *
     * @return 改写用的 SqlSource；无需改写时返回 null
     */
    private SqlSource rewritingSqlSource(MappedStatement ms, Object parameter) {
        SqlNode rootSqlNode = rootSqlNodeOf(ms);
        if (rootSqlNode != null) {
            return new DynamicSqlSource(ms.getConfiguration(),
                    new DataScopeSqlNode(rootSqlNode, ms.getConfiguration(), this));
        }
        BoundSql boundSql = ms.getBoundSql(parameter);
        String newSql = rewriteIfActive(boundSql.getSql());
        if (newSql == null) {
            return null;
        }
        MetaObject metaObject = SystemMetaObject.forObject(boundSql);
        metaObject.setValue("sql", newSql);
        return new FixedBoundSqlSource(boundSql);
    }

    private SqlNode rootSqlNodeOf(MappedStatement ms) {
        if (!(ms.getSqlSource() instanceof DynamicSqlSource)) {
            return null;
        }
        return (SqlNode) SystemMetaObject.forObject(ms).getValue("sqlSource.rootSqlNode");
    }

    private String rewriteIfActive(String sql) {
        if (StringUtils.isBlank(DataScopeHolder.get())) {
            return null;
        }
        return DataScopeSqlRewriter.rewrite(sql, condition, column, tables, dialect);
    }

    /**
     * 把原 mapper 的动态 SQL 包一层：先在临时上下文里求值，再改写 SQL，
     * 并把 &lt;foreach&gt; 等展开值回填到外层绑定。
     */
    private static final class DataScopeSqlNode implements SqlNode {

        private final SqlNode delegate;
        private final Configuration configuration;
        private final DataScopeInterceptor interceptor;

        private DataScopeSqlNode(SqlNode delegate, Configuration configuration, DataScopeInterceptor interceptor) {
            this.delegate = delegate;
            this.configuration = configuration;
            this.interceptor = interceptor;
        }

        @Override
        public boolean apply(DynamicContext context) {
            Object parameterObject = context.getBindings().get(DynamicContext.PARAMETER_OBJECT_KEY);
            DynamicContext scratch = new DynamicContext(configuration, parameterObject);
            if (!delegate.apply(scratch)) {
                return false;
            }
            String sql = scratch.getSql();
            String rewritten = interceptor.rewriteIfActive(sql);
            context.appendSql(rewritten == null ? sql : rewritten);
            for (Map.Entry<String, Object> entry : scratch.getBindings().entrySet()) {
                String key = entry.getKey();
                if (DynamicContext.PARAMETER_OBJECT_KEY.equals(key)
                        || DynamicContext.DATABASE_ID_KEY.equals(key)) {
                    continue;
                }
                context.bind(key, entry.getValue());
            }
            return true;
        }
    }

    private static final class FixedBoundSqlSource implements SqlSource {

        private final BoundSql boundSql;

        private FixedBoundSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }

        @Override
        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }
    }

    /**
     * 用改写后的 BoundSql 构造新 MappedStatement。
     *
     * <p>必须沿用原 id：PagerInterceptor 靠 ms.getId() + "_COUNT" / "_PAGE" 找预定义语句，
     * id 变了分页就找不到。</p>
     */
    private MappedStatement newMappedStatement(MappedStatement ms, SqlSource sqlSource) {
        MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), sqlSource,
                ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.parameterMap(ms.getParameterMap());
        builder.resultMaps(ms.getResultMaps());
        builder.resultSetType(ms.getResultSetType());
        builder.statementType(ms.getStatementType());
        builder.fetchSize(ms.getFetchSize());
        builder.timeout(ms.getTimeout());
        builder.cache(ms.getCache());
        builder.flushCacheRequired(ms.isFlushCacheRequired());
        builder.useCache(ms.isUseCache());
        builder.resultOrdered(ms.isResultOrdered());
        builder.databaseId(ms.getDatabaseId());
        builder.lang(ms.getLang());
        builder.keyGenerator(ms.getKeyGenerator());
        String keyProperties = join(ms.getKeyProperties());
        if (keyProperties != null) {
            builder.keyProperty(keyProperties);
        }
        String keyColumns = join(ms.getKeyColumns());
        if (keyColumns != null) {
            builder.keyColumn(keyColumns);
        }
        return builder.build();
    }

    private static String join(String[] values) {
        if (values == null || values.length == 0) {
            return null;
        }
        StringBuilder builder = new StringBuilder();
        for (String value : values) {
            if (builder.length() > 0) {
                builder.append(",");
            }
            builder.append(value);
        }
        return builder.toString();
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        String tableList = properties.getProperty("tables");
        column = properties.getProperty("column");
        condition = properties.getProperty("condition");
        String configuredDialect = properties.getProperty("dialect");
        if (StringUtils.isNotBlank(configuredDialect)) {
            dialect = configuredDialect;
        }
        if (StringUtils.isBlank(tableList) || StringUtils.isBlank(column) || StringUtils.isBlank(condition)) {
            logger.warn("数据范围过滤插件配置不完整，插件不生效。tables={}, column={}, condition={}",
                    tableList, column, condition);
            enabled = false;
            return;
        }
        if (condition.indexOf('?') >= 0 || condition.indexOf(DataScopeSqlRewriter.COLUMN_PLACEHOLDER) < 0) {
            logger.error("数据范围过滤插件条件模板非法，插件不生效。condition={}", condition);
            enabled = false;
            return;
        }
        Set<String> parsed = new HashSet<String>();
        for (String table : tableList.split(",")) {
            String trimmed = table.trim();
            if (StringUtils.isNotBlank(trimmed)) {
                parsed.add(trimmed.toLowerCase(Locale.ENGLISH));
            }
        }
        if (parsed.isEmpty()) {
            logger.warn("数据范围过滤插件白名单为空，插件不生效。tables={}", tableList);
            enabled = false;
            return;
        }
        tables = parsed;
    }
}

DataScopeSqlRewriter：

package com.common.filter;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLJoinTableSource;
import com.alibaba.druid.sql.ast.statement.SQLSelectQueryBlock;
import com.alibaba.druid.sql.ast.statement.SQLTableSource;
import com.alibaba.druid.sql.visitor.SQLASTVisitorAdapter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Set;

/**
 * 数据范围过滤条件的 SQL 改写器。
 *
 * <p>纯工具类，不依赖 Spring 与 MyBatis，可独立验证。</p>
 *
 * <p>过滤条件以模板形式给出，模板中的 {col} 占位符会被替换成当前 SQL 中该表的限定列引用
 * （优先用别名，无别名时用表名），再由 Druid 解析成表达式挂到 WHERE 上。</p>
 */
public final class DataScopeSqlRewriter {

    private static final Logger logger = LoggerFactory.getLogger(DataScopeSqlRewriter.class);

    /** 条件模板中的列占位符 */
    public static final String COLUMN_PLACEHOLDER = "{col}";

    private DataScopeSqlRewriter() {
    }

    /**
     * 改写 SQL，为白名单表追加过滤条件。
     *
     * @param sql      原始 SQL
     * @param template 条件模板，含 {col} 占位符，如 "{col} != 8"
     * @param column   过滤列名，如 "site_type"
     * @param tables   白名单表名（小写比较）
     * @param dbType   方言名，如 "mariadb"
     * @return 改写后的 SQL；无法安全改写时返回 null，由调用方决定是否放行原 SQL
     */
    public static String rewrite(String sql, String template, String column, Set<String> tables, String dbType) {
        if (sql == null || sql.trim().isEmpty()) {
            return null;
        }
        if (isBlank(template) || isBlank(column) || tables == null || tables.isEmpty()) {
            return null;
        }
        if (template.indexOf('?') >= 0 || template.indexOf(COLUMN_PLACEHOLDER) < 0) {
            logger.error("数据范围过滤条件模板非法，放弃改写。template={}", template);
            return null;
        }
        DbType type = toDbType(dbType);
        try {
            List<SQLStatement> statements = SQLUtils.parseStatements(sql, type);
            if (statements == null || statements.isEmpty()) {
                return null;
            }
            SQLStatement statement = statements.get(0);
            BlockCollector collector = new BlockCollector();
            statement.accept(collector);
            for (SQLSelectQueryBlock block : collector.blocks) {
                List<SQLExpr> whereConditions = new ArrayList<SQLExpr>();
                SQLExpr condition = conditionOf(block.getFrom(), whereConditions, column, tables, template, type);
                for (SQLExpr each : whereConditions) {
                    condition = and(condition, each);
                }
                if (condition != null) {
                    block.setWhere(and(block.getWhere(), condition));
                }
            }
            return SQLUtils.toSQLString(statement, type);
        } catch (Exception e) {
            logger.error("数据范围过滤条件改写失败，该查询未被过滤（fail-open 放行原 SQL），sql={}", sql, e);
            return null;
        }
    }

    /**
     * 收集 SQL 中所有查询块（含子查询、UNION 分支、CTE）。
     *
     * <p>Druid 的 SQLASTVisitorAdapter 默认继续下探，UNION 分支、CTE 与各位置的子查询
     * 都会被 visit(SQLSelectQueryBlock) 收集到，因此无需额外覆写 visit(SQLUnionQuery)。</p>
     */
    private static final class BlockCollector extends SQLASTVisitorAdapter {

        private final List<SQLSelectQueryBlock> blocks = new ArrayList<SQLSelectQueryBlock>();

        @Override
        public boolean visit(SQLSelectQueryBlock x) {
            blocks.add(x);
            return true;
        }
    }

    /**
     * 在表源树中找到白名单表并生成条件。
     *
     * <p>Druid 的 JOIN 树是左深结构：getRight() 恒为叶子表，getLeft() 承载已累积的连接。</p>
     *
     * @param whereConditions 逗号连接时收集到的条件，逗号连接没有 ON 子句，只能挂到 WHERE
     * @return 属于当前作用域（主表）的条件，需要挂到当前块的 WHERE 上
     */
    private static SQLExpr conditionOf(SQLTableSource tableSource, List<SQLExpr> whereConditions,
                                       String column, Set<String> tables, String template, DbType type) {
        if (tableSource instanceof SQLExprTableSource) {
            SQLExprTableSource exprTableSource = (SQLExprTableSource) tableSource;
            String tableName = tableNameOf(exprTableSource);
            if (tableName == null || !tables.contains(tableName.toLowerCase(Locale.ENGLISH))) {
                return null;
            }
            return conditionOf(exprTableSource.getAlias(), tableName, column, template, type);
        }
        if (tableSource instanceof SQLJoinTableSource) {
            SQLJoinTableSource join = (SQLJoinTableSource) tableSource;
            SQLExpr leftCondition = conditionOf(join.getLeft(), whereConditions, column, tables, template, type);
            SQLExpr joinedCondition = joinedConditionOf(join.getRight(), column, tables, template, type);
            if (joinedCondition != null) {
                if (join.getJoinType() == SQLJoinTableSource.JoinType.COMMA) {
                    whereConditions.add(joinedCondition);
                } else if (join.getCondition() == null) {
                    join.setCondition(joinedCondition);
                } else {
                    join.addCondition(joinedCondition);
                }
            }
            return leftCondition;
        }
        return null;
    }

    private static SQLExpr joinedConditionOf(SQLTableSource tableSource, String column, Set<String> tables,
                                             String template, DbType type) {
        if (!(tableSource instanceof SQLExprTableSource)) {
            return null;
        }
        SQLExprTableSource exprTableSource = (SQLExprTableSource) tableSource;
        String tableName = tableNameOf(exprTableSource);
        if (tableName == null || !tables.contains(tableName.toLowerCase(Locale.ENGLISH))) {
            return null;
        }
        return conditionOf(exprTableSource.getAlias(), tableName, column, template, type);
    }

    /** 按限定名生成条件表达式 */
    private static SQLExpr conditionOf(String alias, String tableName, String column, String template, DbType type) {
        String qualifier = (alias == null || alias.trim().isEmpty()) ? tableName : alias;
        String fragment = template.replace(COLUMN_PLACEHOLDER, qualifier + "." + column);
        return SQLUtils.toSQLExpr(fragment, type);
    }

    private static String tableNameOf(SQLExprTableSource tableSource) {
        SQLExpr expr = tableSource.getExpr();
        if (expr instanceof SQLIdentifierExpr) {
            return ((SQLIdentifierExpr) expr).getName();
        }
        if (expr instanceof SQLPropertyExpr) {
            return ((SQLPropertyExpr) expr).getName();
        }
        return null;
    }

    private static SQLExpr and(SQLExpr left, SQLExpr right) {
        if (left == null) {
            return right;
        }
        if (right == null) {
            return left;
        }
        return new SQLBinaryOpExpr(left, right, SQLBinaryOperator.BooleanAnd);
    }

    private static boolean isBlank(String value) {
        return value == null || value.trim().isEmpty();
    }

    private static DbType toDbType(String dbType) {
        if (dbType != null && "mariadb".equalsIgnoreCase(dbType.trim())) {
            return DbType.mariadb;
        }
        return DbType.mysql;
    }
}

向容器内装配过滤器：

@Value("${data.scope.header:}")
private String dataScopeHeader;

/**
 * 数据范围过滤开关 Filter。order=0，必须早于其它 Filter：
 * 开关值要在其它 Filter 触达数据库之前就进 ThreadLocal。
 */
@Bean
public FilterRegistrationBean dataScopeFilterRegistration() {
  FilterRegistrationBean registration = new FilterRegistrationBean();
  registration.setFilter(new DataScopeFilter(dataScopeHeader));
  registration.addUrlPatterns("/*");
  registration.setName("dataScopeFilter");
  registration.setOrder(0);
  return registration;
}

修改mybatis-config.xml配置：在plugins中添加：

<!--
    数据范围过滤。必须声明在 PagerInterceptor 之前（即更内层）：
    MyBatis 的 <plugins> 顺序是反的，声明在最后的先执行。
    PagerInterceptor 内部执行分页的 COUNT 查询时走 invocation.proceed()，
    目标是 Executor；若本插件在最外层，COUNT 查询会绕过它，
    导致「列表被过滤、count 没被过滤」，总数与列表口径不一致。
    放在更内层才能拦到 Pager 产出的 count 与分页两条语句。
    mapper 里预定义了 xxx_COUNT 语句时尤其要注意这一点。
    条件模板里的 {col} 会被替换成「别名.列名」，不要写成 ${col}：
    src/main/resources 开了 Maven 资源过滤，${...} 会在打包期被替换掉。
-->
<plugin interceptor="com.stock.capital.nmas.common.filter.DataScopeInterceptor">
<property name="tables" value="table_name1,table_name2"/>
    <property name="column" value="column_name1"/>
    <!-- value为添加的过滤条件表达式，{col}固定不变 --> 
    <property name="condition" value="{col} != 8"/>
    <property name="dialect" value="mariadb"/>
</plugin>

bootstrap.yml中添加：

# 数据范围过滤开关：请求头名。有值则对白名单表追加过滤条件，无值则不过滤
data.scope.header=X-Data-Scope

对应的请求头通过Nginx配置添加，至此，业务代码不需要任何改动，以最小改动完美实现。✌️