Implementing Data Permission Control with MyBatis Interceptor and ThreadLocal

1. Usage

Data permission control requires filtering the data returned by queries. The least invasive approach is to modify existing SQL through MyBatis or database‑pool interceptors, adding only a small amount of marker code in the business layer. This method leaves legacy code almost untouched and lets new development focus on core logic.

Because many different connection‑pool implementations are used across departments, a unified slicing solution is difficult. Most SQL is built either via MyBatis or by direct Java string concatenation, so the solution adopts MyBatis slicing. By adding an annotation to a mapper interface that specifies the permission type, table name, and column name, the required data permissions are automatically applied.

Since the same mapper method may be invoked in contexts that need permission control and others that do not, an additional mechanism is introduced: passing permission rules through a ThreadLocal so that the current SQL execution can be governed accordingly.

Permission rule format example:

限权规则code1(表名1.字段名1,表名2.字段名2);限权规则code2(表名3.字段名3,表名4.字段名4)

Example usage:

enterprise(channel.enterprise_code);account(table.column);channel(table3.id)

The context‑passing utility class uses callbacks to ensure the ThreadLocal is cleared after use:

public class DataAuthContextUtil {
private static ThreadLocal<String> useDataAuth = new ThreadLocal<>();
public static <T> T executeSqlWithDataAuthRule(String rule, Supplier<T> supplier) {
try {
useDataAuth.set(rule);
return supplier.get();
} finally {
useDataAuth.remove();
}
}
public static String getUseDataAuthRule() {
return useDataAuth.get();
}
}

2. Slice Implementation Process

3. Other Technical Details

(1) Retrieve the original SQL in the interceptor

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.cache.CacheKey;
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.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.util.function.Tuple2;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
@Component
@Intercepts({
@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
@Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
})
@Slf4j
public class DataAuthInterceptor implements Interceptor {
@Override
public Object intercept(Invocation invocation) throws Throwable {
try {
MappedStatement mappedStatement = (MappedStatement) invocation.getArgs()[0];
BoundSql boundSql = mappedStatement.getBoundSql(invocation.getArgs()[1]);
String sql = boundSql.getSql();
// permission logic will modify sql here
} catch (Exception e) {
log.error("Data permission addition error, current sql has no permission restriction!", e);
throw e;
}
return invocation.proceed();
}
}

(2) Add permission items into the original SQL

The Druid AST parser is used to modify the SQL as shown below:

/**
* Write permission limits into SQL
* @param sql original SQL
* @param tableAuthMap key:table, value:column and values for permission
*/
public static StringBuilder addAuthLimitToSql(String sql, Map<String, Tuple2<String, Set<String>>> tableAuthMap) {
List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, "mysql");
StringBuilder authSql = new StringBuilder();
for (SQLStatement stmt : stmtList) {
stmt.accept(new MySqlASTVisitorAdapter() {
@Override
public boolean visit(MySqlSelectQueryBlock x) {
SQLTableSource from = x.getFrom();
Set<String> tableList = new HashSet<>();
getTableList(from, tableList);
for (String tableName : tableList) {
if (tableAuthMap.containsKey(tableName)) {
x.addCondition(tableName + "in (...略)");
}
}
return true;
}
});
authSql.append(stmt);
}
return authSql;
}
private static void getTableList(SQLTableSource from, Set<String> tableList) {
if (from instanceof SQLExprTableSource) {
SQLExprTableSource tableSource = (SQLExprTableSource) from;
String name = tableSource.getTableName().replace("`", "");
tableList.add(name);
String alias = tableSource.getAlias();
if (StringUtils.isNotBlank(alias)) {
tableList.add(alias.replace("`", ""));
}
} else if (from instanceof SQLJoinTableSource) {
SQLJoinTableSource join = (SQLJoinTableSource) from;
getTableList(join.getLeft(), tableList);
getTableList(join.getRight(), tableList);
} else if (from instanceof SQLSubqueryTableSource) {
SQLSubqueryTableSource sub = (SQLSubqueryTableSource) from;
tableList.add(sub.getAlias().replace("`", ""));
} else {
// other source types are ignored or logged
}
}

(3) Write the modified SQL back to MyBatis

MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
BoundSql boundSql = ms.getBoundSql(invocation.getArgs()[1]);
// Build a new MappedStatement with the modified SQL
MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), new MySqlSource(boundSql), ms.getSqlCommandType());
builder.resource(ms.getResource());
builder.fetchSize(ms.getFetchSize());
builder.statementType(ms.getStatementType());
builder.keyGenerator(ms.getKeyGenerator());
if (ms.getKeyProperties() != null && ms.getKeyProperties().length != 0) {
StringBuilder keyProperties = new StringBuilder();
for (String keyProperty : ms.getKeyProperties()) {
keyProperties.append(keyProperty).append(",");
}
keyProperties.delete(keyProperties.length() - 1, keyProperties.length());
builder.keyProperty(keyProperties.toString());
}
builder.timeout(ms.getTimeout());
builder.parameterMap(ms.getParameterMap());
builder.resultMaps(ms.getResultMaps());
builder.resultSetType(ms.getResultSetType());
builder.cache(ms.getCache());
builder.flushCacheRequired(ms.isFlushCacheRequired());
builder.useCache(ms.isUseCache());
MappedStatement newMappedStatement = builder.build();
MetaObject metaObject = MetaObject.forObject(newMappedStatement, new DefaultObjectFactory(), new DefaultObjectWrapperFactory(), new DefaultReflectorFactory());
metaObject.setValue("sqlSource.boundSql.sql", newSql);
invocation.getArgs()[0] = newMappedStatement;

Reference article: https://blog.csdn.net/e_anjing/article/details/79102693