BJ_ZTEITS-Developers / parking-cloud-yudao

Commit bea66e5c6f1f2582c405411645ac30935a2f9b56

Authored by YunaiV
1 parent 2b39b434

【代码优化】全局：复用 MyBatis Plus 数据权限插件，简化项目的数据权限实现

Inline Side-by-side
Showing 5 changed files with 77 additions and 499 deletions
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/config/YudaoDataPermissionAutoConfiguration.java
  View file @bea66e5
1 1 package cn.iocoder.yudao.framework.datapermission.config;
2 2  
3 3 import cn.iocoder.yudao.framework.datapermission.core.aop.DataPermissionAnnotationAdvisor;
4   -import cn.iocoder.yudao.framework.datapermission.core.db.DataPermissionDatabaseInterceptor;
  4 +import cn.iocoder.yudao.framework.datapermission.core.db.DataPermissionRuleHandler;
5 5 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
6 6 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
7 7 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactoryImpl;
8 8 import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
9 9 import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
  10 +import com.baomidou.mybatisplus.extension.plugins.inner.DataPermissionInterceptor;
10 11 import org.springframework.boot.autoconfigure.AutoConfiguration;
11 12 import org.springframework.context.annotation.Bean;
12 13  
... ... @@ -26,14 +27,15 @@ public class YudaoDataPermissionAutoConfiguration {
26 27 }
27 28  
28 29 @Bean
29   - public DataPermissionDatabaseInterceptor dataPermissionDatabaseInterceptor(MybatisPlusInterceptor interceptor,
30   - DataPermissionRuleFactory ruleFactory) {
31   - // 创建 DataPermissionDatabaseInterceptor 拦截器
32   - DataPermissionDatabaseInterceptor inner = new DataPermissionDatabaseInterceptor(ruleFactory);
  30 + public DataPermissionRuleHandler dataPermissionRuleHandler(MybatisPlusInterceptor interceptor,
  31 + DataPermissionRuleFactory ruleFactory) {
  32 + // 创建 DataPermissionInterceptor 拦截器
  33 + DataPermissionRuleHandler handler = new DataPermissionRuleHandler(ruleFactory);
  34 + DataPermissionInterceptor inner = new DataPermissionInterceptor(handler);
33 35 // 添加到 interceptor 中
34 36 // 需要加在首个，主要是为了在分页插件前面。这个是 MyBatis Plus 的规定
35 37 MyBatisUtils.addInterceptor(interceptor, inner, 0);
36   - return inner;
  38 + return handler;
37 39 }
38 40  
39 41 @Bean
... ...
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/core/db/DataPermissionDatabaseInterceptor.java deleted
View file @2b39b43
1   -package cn.iocoder.yudao.framework.datapermission.core.db;
2   -
3   -import cn.hutool.core.collection.CollUtil;
4   -import cn.iocoder.yudao.framework.common.util.collection.SetUtils;
5   -import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
6   -import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
7   -import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
8   -import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
9   -import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
10   -import com.baomidou.mybatisplus.extension.plugins.inner.BaseMultiTableInnerInterceptor;
11   -import com.baomidou.mybatisplus.extension.plugins.inner.DataPermissionInterceptor;
12   -import com.baomidou.mybatisplus.extension.plugins.inner.InnerInterceptor;
13   -import lombok.Getter;
14   -import lombok.RequiredArgsConstructor;
15   -import net.sf.jsqlparser.expression.Expression;
16   -import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
17   -import net.sf.jsqlparser.schema.Table;
18   -import net.sf.jsqlparser.statement.delete.Delete;
19   -import net.sf.jsqlparser.statement.select.Select;
20   -import net.sf.jsqlparser.statement.select.WithItem;
21   -import net.sf.jsqlparser.statement.update.Update;
22   -import org.apache.ibatis.executor.Executor;
23   -import org.apache.ibatis.executor.statement.StatementHandler;
24   -import org.apache.ibatis.mapping.BoundSql;
25   -import org.apache.ibatis.mapping.MappedStatement;
26   -import org.apache.ibatis.mapping.SqlCommandType;
27   -import org.apache.ibatis.session.ResultHandler;
28   -import org.apache.ibatis.session.RowBounds;
29   -
30   -import java.sql.Connection;
31   -import java.util.Collections;
32   -import java.util.List;
33   -import java.util.Map;
34   -import java.util.Set;
35   -import java.util.concurrent.ConcurrentHashMap;
36   -
37   -/**
38   - * 数据权限拦截器，通过 {@link DataPermissionRule} 数据权限规则，重写 SQL 的方式来实现
39   - * 主要的 SQL 重写方法，可见 {@link #buildTableExpression(Table, Expression, String)} 方法
40   - *
41   - * 整体的代码实现上，参考 {@link com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor} 实现。
42   - * 所以每次 MyBatis Plus 升级时，需要 Review 下其具体的实现是否有变更！
43   - *
44   - * 为什么不直接基于 {@link DataPermissionInterceptor} 实现？因为它是 MyBatis Plus 3.5.2 版本才出来，当时 yudao 已经实现数据权限了！
45   - *
46   - * @author 芋道源码
47   - */
48   -@RequiredArgsConstructor
49   -public class DataPermissionDatabaseInterceptor extends BaseMultiTableInnerInterceptor implements InnerInterceptor {
50   -
51   - private final DataPermissionRuleFactory ruleFactory;
52   -
53   - @Getter
54   - private final MappedStatementCache mappedStatementCache = new MappedStatementCache();
55   -
56   - @Override // SELECT 场景
57   - public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
58   - // 获得 Mapper 对应的数据权限的规则
59   - List<DataPermissionRule> rules = ruleFactory.getDataPermissionRule(ms.getId());
60   - if (mappedStatementCache.noRewritable(ms, rules)) { // 如果无需重写，则跳过
61   - return;
62   - }
63   -
64   - PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
65   - try {
66   - // 初始化上下文
67   - ContextHolder.init(rules);
68   - // 处理 SQL
69   - mpBs.sql(parserSingle(mpBs.sql(), null));
70   - } finally {
71   - // 添加是否需要重写的缓存
72   - addMappedStatementCache(ms);
73   - // 清空上下文
74   - ContextHolder.clear();
75   - }
76   - }
77   -
78   - @Override // 只处理 UPDATE / DELETE 场景，不处理 INSERT 场景（因为 INSERT 不需要数据权限)
79   - public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
80   - PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
81   - MappedStatement ms = mpSh.mappedStatement();
82   - SqlCommandType sct = ms.getSqlCommandType();
83   - if (sct == SqlCommandType.UPDATE || sct == SqlCommandType.DELETE) {
84   - // 获得 Mapper 对应的数据权限的规则
85   - List<DataPermissionRule> rules = ruleFactory.getDataPermissionRule(ms.getId());
86   - if (mappedStatementCache.noRewritable(ms, rules)) { // 如果无需重写，则跳过
87   - return;
88   - }
89   -
90   - PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql();
91   - try {
92   - // 初始化上下文
93   - ContextHolder.init(rules);
94   - // 处理 SQL
95   - mpBs.sql(parserMulti(mpBs.sql(), null));
96   - } finally {
97   - // 添加是否需要重写的缓存
98   - addMappedStatementCache(ms);
99   - // 清空上下文
100   - ContextHolder.clear();
101   - }
102   - }
103   - }
104   -
105   - @Override
106   - protected void processSelect(Select select, int index, String sql, Object obj) {
107   - final String whereSegment = (String) obj;
108   - processSelectBody(select, whereSegment);
109   - List<WithItem> withItemsList = select.getWithItemsList();
110   - if (!CollectionUtils.isEmpty(withItemsList)) {
111   - withItemsList.forEach(withItem -> processSelectBody(withItem, whereSegment));
112   - }
113   - }
114   -
115   - /**
116   - * update 语句处理
117   - */
118   - @Override
119   - protected void processUpdate(Update update, int index, String sql, Object obj) {
120   - final Expression sqlSegment = getUpdateOrDeleteExpression(update.getTable(), update.getWhere(), (String) obj);
121   - if (null != sqlSegment) {
122   - update.setWhere(sqlSegment);
123   - }
124   - }
125   -
126   - /**
127   - * delete 语句处理
128   - */
129   - @Override
130   - protected void processDelete(Delete delete, int index, String sql, Object obj) {
131   - final Expression sqlSegment = getUpdateOrDeleteExpression(delete.getTable(), delete.getWhere(), (String) obj);
132   - if (null != sqlSegment) {
133   - delete.setWhere(sqlSegment);
134   - }
135   - }
136   -
137   - // ========== 和 TenantLineInnerInterceptor 存在差异的逻辑：关键，实现权限条件的拼接 ==========
138   -
139   - protected Expression getUpdateOrDeleteExpression(final Table table, final Expression where, final String whereSegment) {
140   - return andExpression(table, where, whereSegment);
141   - }
142   -
143   - /**
144   - * 构建指定表的数据权限的 Expression 过滤条件
145   - *
146   - * @param table 表
147   - * @return Expression 过滤条件
148   - */
149   - @Override
150   - public Expression buildTableExpression(Table table, Expression where, String whereSegment) {
151   - // 生成条件
152   - Expression allExpression = null;
153   - for (DataPermissionRule rule : ContextHolder.getRules()) {
154   - // 判断表名是否匹配
155   - String tableName = MyBatisUtils.getTableName(table);
156   - if (!rule.getTableNames().contains(tableName)) {
157   - continue;
158   - }
159   - // 如果有匹配的规则，说明可重写。
160   - // 为什么不是有 allExpression 非空才重写呢？在生成 column = value 过滤条件时，会因为 value 不存在，导致未重写。
161   - // 这样导致第一次无 value，被标记成无需重写；但是第二次有 value，此时会需要重写。
162   - ContextHolder.setRewrite(true);
163   -
164   - // 单条规则的条件
165   - Expression oneExpress = rule.getExpression(tableName, table.getAlias());
166   - if (oneExpress == null){
167   - continue;
168   - }
169   - // 拼接到 allExpression 中
170   - allExpression = allExpression == null ? oneExpress
171   - : new AndExpression(allExpression, oneExpress);
172   - }
173   -
174   - return allExpression;
175   - }
176   -
177   - /**
178   - * 判断 SQL 是否重写。如果没有重写，则添加到 {@link MappedStatementCache} 中
179   - *
180   - * @param ms MappedStatement
181   - */
182   - private void addMappedStatementCache(MappedStatement ms) {
183   - if (ContextHolder.getRewrite()) {
184   - return;
185   - }
186   - // 无重写，进行添加
187   - mappedStatementCache.addNoRewritable(ms, ContextHolder.getRules());
188   - }
189   -
190   - /**
191   - * SQL 解析上下文，方便透传 {@link DataPermissionRule} 规则
192   - *
193   - * @author 芋道源码
194   - */
195   - static final class ContextHolder {
196   -
197   - /**
198   - * 该 {@link MappedStatement} 对应的规则
199   - */
200   - private static final ThreadLocal<List<DataPermissionRule>> RULES = ThreadLocal.withInitial(Collections::emptyList);
201   - /**
202   - * SQL 是否进行重写
203   - */
204   - private static final ThreadLocal<Boolean> REWRITE = ThreadLocal.withInitial(() -> Boolean.FALSE);
205   -
206   - public static void init(List<DataPermissionRule> rules) {
207   - RULES.set(rules);
208   - REWRITE.set(false);
209   - }
210   -
211   - public static void clear() {
212   - RULES.remove();
213   - REWRITE.remove();
214   - }
215   -
216   - public static boolean getRewrite() {
217   - return REWRITE.get();
218   - }
219   -
220   - public static void setRewrite(boolean rewrite) {
221   - REWRITE.set(rewrite);
222   - }
223   -
224   - public static List<DataPermissionRule> getRules() {
225   - return RULES.get();
226   - }
227   -
228   - }
229   -
230   - /**
231   - * {@link MappedStatement} 缓存
232   - * 目前主要用于，记录 {@link DataPermissionRule} 是否对指定 {@link MappedStatement} 无效
233   - * 如果无效，则可以避免 SQL 的解析，加快速度
234   - *
235   - * @author 芋道源码
236   - */
237   - static final class MappedStatementCache {
238   -
239   - /**
240   - * 指定数据权限规则，对指定 MappedStatement 无需重写（不生效)的缓存
241   - *
242   - * value：{@link MappedStatement#getId()} 编号
243   - */
244   - @Getter
245   - private final Map<Class<? extends DataPermissionRule>, Set<String>> noRewritableMappedStatements = new ConcurrentHashMap<>();
246   -
247   - /**
248   - * 判断是否无需重写
249   - * ps：虽然有点中文式英语，但是容易读懂即可
250   - *
251   - * @param ms MappedStatement
252   - * @param rules 数据权限规则数组
253   - * @return 是否无需重写
254   - */
255   - public boolean noRewritable(MappedStatement ms, List<DataPermissionRule> rules) {
256   - // 如果规则为空，说明无需重写
257   - if (CollUtil.isEmpty(rules)) {
258   - return true;
259   - }
260   - // 任一规则不在 noRewritableMap 中，则说明可能需要重写
261   - for (DataPermissionRule rule : rules) {
262   - Set<String> mappedStatementIds = noRewritableMappedStatements.get(rule.getClass());
263   - if (!CollUtil.contains(mappedStatementIds, ms.getId())) {
264   - return false;
265   - }
266   - }
267   - return true;
268   - }
269   -
270   - /**
271   - * 添加无需重写的 MappedStatement
272   - *
273   - * @param ms MappedStatement
274   - * @param rules 数据权限规则数组
275   - */
276   - public void addNoRewritable(MappedStatement ms, List<DataPermissionRule> rules) {
277   - for (DataPermissionRule rule : rules) {
278   - Set<String> mappedStatementIds = noRewritableMappedStatements.get(rule.getClass());
279   - if (CollUtil.isNotEmpty(mappedStatementIds)) {
280   - mappedStatementIds.add(ms.getId());
281   - } else {
282   - noRewritableMappedStatements.put(rule.getClass(), SetUtils.asSet(ms.getId()));
283   - }
284   - }
285   - }
286   -
287   - /**
288   - * 清空缓存
289   - * 目前主要提供给单元测试
290   - */
291   - public void clear() {
292   - noRewritableMappedStatements.clear();
293   - }
294   -
295   - }
296   -
297   -}
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/core/db/DataPermissionRuleHandler.java 0 → 100644
  View file @bea66e5
  1 +package cn.iocoder.yudao.framework.datapermission.core.db;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
  5 +import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
  6 +import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
  7 +import com.baomidou.mybatisplus.extension.plugins.handler.MultiDataPermissionHandler;
  8 +import lombok.RequiredArgsConstructor;
  9 +import net.sf.jsqlparser.expression.Expression;
  10 +import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
  11 +import net.sf.jsqlparser.schema.Table;
  12 +
  13 +import java.util.List;
  14 +
  15 +/**
  16 + * 基于 {@link DataPermissionRule} 的数据权限处理器
  17 + *
  18 + * 它的底层，是基于 MyBatis Plus 的 <a href="https://baomidou.com/plugins/data-permission/">数据权限插件</a>
  19 + * 核心原理：它会在 SQL 执行前拦截 SQL 语句，并根据用户权限动态添加权限相关的 SQL 片段。这样，只有用户有权限访问的数据才会被查询出来
  20 + *
  21 + * @author 芋道源码
  22 + */
  23 +@RequiredArgsConstructor
  24 +public class DataPermissionRuleHandler implements MultiDataPermissionHandler {
  25 +
  26 + private final DataPermissionRuleFactory ruleFactory;
  27 +
  28 + @Override
  29 + public Expression getSqlSegment(Table table, Expression where, String mappedStatementId) {
  30 + // 获得 Mapper 对应的数据权限的规则
  31 + List<DataPermissionRule> rules = ruleFactory.getDataPermissionRule(mappedStatementId);
  32 + if (CollUtil.isEmpty(rules)) {
  33 + return null;
  34 + }
  35 +
  36 + // 生成条件
  37 + Expression allExpression = null;
  38 + for (DataPermissionRule rule : rules) {
  39 + // 判断表名是否匹配
  40 + String tableName = MyBatisUtils.getTableName(table);
  41 + if (!rule.getTableNames().contains(tableName)) {
  42 + continue;
  43 + }
  44 +
  45 + // 单条规则的条件
  46 + Expression oneExpress = rule.getExpression(tableName, table.getAlias());
  47 + if (oneExpress == null) {
  48 + continue;
  49 + }
  50 + // 拼接到 allExpression 中
  51 + allExpression = allExpression == null ? oneExpress
  52 + : new AndExpression(allExpression, oneExpress);
  53 + }
  54 + return allExpression;
  55 + }
  56 +
  57 +}
... ...
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/test/java/cn/iocoder/yudao/framework/datapermission/core/db/DataPermissionDatabaseInterceptorTest.java deleted
View file @2b39b43
1   -package cn.iocoder.yudao.framework.datapermission.core.db;
2   -
3   -import cn.iocoder.yudao.framework.common.util.collection.SetUtils;
4   -import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
5   -import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
6   -import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
7   -import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
8   -import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
9   -import net.sf.jsqlparser.expression.Alias;
10   -import net.sf.jsqlparser.expression.Expression;
11   -import net.sf.jsqlparser.expression.LongValue;
12   -import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
13   -import net.sf.jsqlparser.schema.Column;
14   -import org.apache.ibatis.executor.Executor;
15   -import org.apache.ibatis.executor.statement.StatementHandler;
16   -import org.apache.ibatis.mapping.BoundSql;
17   -import org.apache.ibatis.mapping.MappedStatement;
18   -import org.junit.jupiter.api.BeforeEach;
19   -import org.junit.jupiter.api.Test;
20   -import org.mockito.InjectMocks;
21   -import org.mockito.Mock;
22   -import org.mockito.MockedStatic;
23   -
24   -import java.sql.Connection;
25   -import java.util.*;
26   -
27   -import static java.util.Collections.singletonList;
28   -import static org.junit.jupiter.api.Assertions.*;
29   -import static org.mockito.Mockito.*;
30   -
31   -/**
32   - * {@link DataPermissionDatabaseInterceptor} 的单元测试
33   - * 主要测试 {@link DataPermissionDatabaseInterceptor#beforePrepare(StatementHandler, Connection, Integer)}
34   - * 和 {@link DataPermissionDatabaseInterceptor#beforeUpdate(Executor, MappedStatement, Object)}
35   - * 以及在这个过程中，ContextHolder 和 MappedStatementCache
36   - *
37   - * @author 芋道源码
38   - */
39   -public class DataPermissionDatabaseInterceptorTest extends BaseMockitoUnitTest {
40   -
41   - @InjectMocks
42   - private DataPermissionDatabaseInterceptor interceptor;
43   -
44   - @Mock
45   - private DataPermissionRuleFactory ruleFactory;
46   -
47   - @BeforeEach
48   - public void setUp() {
49   - // 清理上下文
50   - DataPermissionDatabaseInterceptor.ContextHolder.clear();
51   - // 清空缓存
52   - interceptor.getMappedStatementCache().clear();
53   - }
54   -
55   - @Test // 不存在规则，且不匹配
56   - public void testBeforeQuery_withoutRule() {
57   - try (MockedStatic<PluginUtils> pluginUtilsMock = mockStatic(PluginUtils.class)) {
58   - // 准备参数
59   - MappedStatement mappedStatement = mock(MappedStatement.class);
60   - BoundSql boundSql = mock(BoundSql.class);
61   -
62   - // 调用
63   - interceptor.beforeQuery(null, mappedStatement, null, null, null, boundSql);
64   - // 断言
65   - pluginUtilsMock.verify(() -> PluginUtils.mpBoundSql(boundSql), never());
66   - }
67   - }
68   -
69   - @Test // 存在规则，且不匹配
70   - public void testBeforeQuery_withMatchRule() {
71   - try (MockedStatic<PluginUtils> pluginUtilsMock = mockStatic(PluginUtils.class)) {
72   - // 准备参数
73   - MappedStatement mappedStatement = mock(MappedStatement.class);
74   - BoundSql boundSql = mock(BoundSql.class);
75   - // mock 方法(数据权限)
76   - when(ruleFactory.getDataPermissionRule(same(mappedStatement.getId())))
77   - .thenReturn(singletonList(new DeptDataPermissionRule()));
78   - // mock 方法(MPBoundSql)
79   - PluginUtils.MPBoundSql mpBs = mock(PluginUtils.MPBoundSql.class);
80   - pluginUtilsMock.when(() -> PluginUtils.mpBoundSql(same(boundSql))).thenReturn(mpBs);
81   - // mock 方法(SQL)
82   - String sql = "select * from t_user where id = 1";
83   - when(mpBs.sql()).thenReturn(sql);
84   - // 针对 ContextHolder 和 MappedStatementCache 暂时不 mock，主要想校验过程中，数据是否正确
85   -
86   - // 调用
87   - interceptor.beforeQuery(null, mappedStatement, null, null, null, boundSql);
88   - // 断言
89   - verify(mpBs, times(1)).sql(
90   - eq("SELECT * FROM t_user WHERE id = 1 AND t_user.dept_id = 100"));
91   - // 断言缓存
92   - assertTrue(interceptor.getMappedStatementCache().getNoRewritableMappedStatements().isEmpty());
93   - }
94   - }
95   -
96   - @Test // 存在规则，但不匹配
97   - public void testBeforeQuery_withoutMatchRule() {
98   - try (MockedStatic<PluginUtils> pluginUtilsMock = mockStatic(PluginUtils.class)) {
99   - // 准备参数
100   - MappedStatement mappedStatement = mock(MappedStatement.class);
101   - BoundSql boundSql = mock(BoundSql.class);
102   - // mock 方法(数据权限)
103   - when(ruleFactory.getDataPermissionRule(same(mappedStatement.getId())))
104   - .thenReturn(singletonList(new DeptDataPermissionRule()));
105   - // mock 方法(MPBoundSql)
106   - PluginUtils.MPBoundSql mpBs = mock(PluginUtils.MPBoundSql.class);
107   - pluginUtilsMock.when(() -> PluginUtils.mpBoundSql(same(boundSql))).thenReturn(mpBs);
108   - // mock 方法(SQL)
109   - String sql = "select * from t_role where id = 1";
110   - when(mpBs.sql()).thenReturn(sql);
111   - // 针对 ContextHolder 和 MappedStatementCache 暂时不 mock，主要想校验过程中，数据是否正确
112   -
113   - // 调用
114   - interceptor.beforeQuery(null, mappedStatement, null, null, null, boundSql);
115   - // 断言
116   - verify(mpBs, times(1)).sql(
117   - eq("SELECT * FROM t_role WHERE id = 1"));
118   - // 断言缓存
119   - assertFalse(interceptor.getMappedStatementCache().getNoRewritableMappedStatements().isEmpty());
120   - }
121   - }
122   -
123   - @Test
124   - public void testAddNoRewritable() {
125   - // 准备参数
126   - MappedStatement ms = mock(MappedStatement.class);
127   - List<DataPermissionRule> rules = singletonList(new DeptDataPermissionRule());
128   - // mock 方法
129   - when(ms.getId()).thenReturn("selectById");
130   -
131   - // 调用
132   - interceptor.getMappedStatementCache().addNoRewritable(ms, rules);
133   - // 断言
134   - Map<Class<? extends DataPermissionRule>, Set<String>> noRewritableMappedStatements =
135   - interceptor.getMappedStatementCache().getNoRewritableMappedStatements();
136   - assertEquals(1, noRewritableMappedStatements.size());
137   - assertEquals(SetUtils.asSet("selectById"), noRewritableMappedStatements.get(DeptDataPermissionRule.class));
138   - }
139   -
140   - @Test
141   - public void testNoRewritable() {
142   - // 准备参数
143   - MappedStatement ms = mock(MappedStatement.class);
144   - // mock 方法
145   - when(ms.getId()).thenReturn("selectById");
146   - // mock 数据
147   - List<DataPermissionRule> rules = singletonList(new DeptDataPermissionRule());
148   - interceptor.getMappedStatementCache().addNoRewritable(ms, rules);
149   -
150   - // 场景一，rules 为空
151   - assertTrue(interceptor.getMappedStatementCache().noRewritable(ms, null));
152   - // 场景二，rules 非空，可重写
153   - assertFalse(interceptor.getMappedStatementCache().noRewritable(ms, singletonList(new EmptyDataPermissionRule())));
154   - // 场景三，rule 非空，不可重写
155   - assertTrue(interceptor.getMappedStatementCache().noRewritable(ms, rules));
156   - }
157   -
158   - private static class DeptDataPermissionRule implements DataPermissionRule {
159   -
160   - private static final String COLUMN = "dept_id";
161   -
162   - @Override
163   - public Set<String> getTableNames() {
164   - return SetUtils.asSet("t_user");
165   - }
166   -
167   - @Override
168   - public Expression getExpression(String tableName, Alias tableAlias) {
169   - Column column = MyBatisUtils.buildColumn(tableName, tableAlias, COLUMN);
170   - LongValue value = new LongValue(100L);
171   - return new EqualsTo(column, value);
172   - }
173   -
174   - }
175   -
176   - private static class EmptyDataPermissionRule implements DataPermissionRule {
177   -
178   - @Override
179   - public Set<String> getTableNames() {
180   - return Collections.emptySet();
181   - }
182   -
183   - @Override
184   - public Expression getExpression(String tableName, Alias tableAlias) {
185   - return null;
186   - }
187   -
188   - }
189   -
190   -}
yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/test/java/cn/iocoder/yudao/framework/datapermission/core/db/DataPermissionDatabaseInterceptorTest2.java renamed to yudao-framework/yudao-spring-boot-starter-biz-data-permission/src/test/java/cn/iocoder/yudao/framework/datapermission/core/db/DataPermissionRuleHandlerTest.java
  View file @bea66e5
... ... @@ -4,6 +4,7 @@ import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
4 4 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
5 5 import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
6 6 import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  7 +import com.baomidou.mybatisplus.extension.plugins.inner.DataPermissionInterceptor;
7 8 import net.sf.jsqlparser.expression.Alias;
8 9 import net.sf.jsqlparser.expression.Expression;
9 10 import net.sf.jsqlparser.expression.LongValue;
... ... @@ -22,24 +23,30 @@ import java.util.Set;
22 23  
23 24 import static cn.iocoder.yudao.framework.common.util.collection.SetUtils.asSet;
24 25 import static org.junit.jupiter.api.Assertions.assertEquals;
  26 +import static org.mockito.ArgumentMatchers.any;
  27 +import static org.mockito.Mockito.when;
25 28  
26 29 /**
27   - * {@link DataPermissionDatabaseInterceptor} 的单元测试
  30 + * {@link DataPermissionRuleHandler} 的单元测试
28 31 * 主要复用了 MyBatis Plus 的 TenantLineInnerInterceptorTest 的单元测试
29 32 * 不过它的单元测试不是很规范，考虑到是复用的，所以暂时不进行修改~
30 33 *
31 34 * @author 芋道源码
32 35 */
33   -public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest {
  36 +public class DataPermissionRuleHandlerTest extends BaseMockitoUnitTest {
34 37  
35 38 @InjectMocks
36   - private DataPermissionDatabaseInterceptor interceptor;
  39 + private DataPermissionRuleHandler handler;
37 40  
38 41 @Mock
39 42 private DataPermissionRuleFactory ruleFactory;
40 43  
  44 + private DataPermissionInterceptor interceptor;
  45 +
41 46 @BeforeEach
42 47 public void setUp() {
  48 + interceptor = new DataPermissionInterceptor(handler);
  49 +
43 50 // 租户的数据权限规则
44 51 DataPermissionRule tenantRule = new DataPermissionRule() {
45 52  
... ... @@ -78,8 +85,8 @@ public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest
78 85 }
79 86  
80 87 };
81   - // 设置到上下文，保证
82   - DataPermissionDatabaseInterceptor.ContextHolder.init(Arrays.asList(tenantRule, deptRule));
  88 + // 设置到上下文
  89 + when(ruleFactory.getDataPermissionRule(any())).thenReturn(Arrays.asList(tenantRule, deptRule));
83 90 }
84 91  
85 92 @Test
... ... @@ -448,7 +455,6 @@ public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest
448 455 assertEquals(targetSql, interceptor.parserSingle(sql, null));
449 456 }
450 457  
451   -
452 458 // ========== 额外的测试 ==========
453 459  
454 460 @Test
... ...