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 package cn.iocoder.yudao.framework.datapermission.config; 1 package cn.iocoder.yudao.framework.datapermission.config;
2 2
3 import cn.iocoder.yudao.framework.datapermission.core.aop.DataPermissionAnnotationAdvisor; 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 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule; 5 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
6 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory; 6 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
7 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactoryImpl; 7 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactoryImpl;
8 import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils; 8 import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
9 import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor; 9 import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
  10 +import com.baomidou.mybatisplus.extension.plugins.inner.DataPermissionInterceptor;
10 import org.springframework.boot.autoconfigure.AutoConfiguration; 11 import org.springframework.boot.autoconfigure.AutoConfiguration;
11 import org.springframework.context.annotation.Bean; 12 import org.springframework.context.annotation.Bean;
12 13
@@ -26,14 +27,15 @@ public class YudaoDataPermissionAutoConfiguration { @@ -26,14 +27,15 @@ public class YudaoDataPermissionAutoConfiguration {
26 } 27 }
27 28
28 @Bean 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 // 添加到 interceptor 中 35 // 添加到 interceptor 中
34 // 需要加在首个，主要是为了在分页插件前面。这个是 MyBatis Plus 的规定 36 // 需要加在首个，主要是为了在分页插件前面。这个是 MyBatis Plus 的规定
35 MyBatisUtils.addInterceptor(interceptor, inner, 0); 37 MyBatisUtils.addInterceptor(interceptor, inner, 0);
36 - return inner; 38 + return handler;
37 } 39 }
38 40
39 @Bean 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,6 +4,7 @@ import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRule;
4 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory; 4 import cn.iocoder.yudao.framework.datapermission.core.rule.DataPermissionRuleFactory;
5 import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils; 5 import cn.iocoder.yudao.framework.mybatis.core.util.MyBatisUtils;
6 import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest; 6 import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  7 +import com.baomidou.mybatisplus.extension.plugins.inner.DataPermissionInterceptor;
7 import net.sf.jsqlparser.expression.Alias; 8 import net.sf.jsqlparser.expression.Alias;
8 import net.sf.jsqlparser.expression.Expression; 9 import net.sf.jsqlparser.expression.Expression;
9 import net.sf.jsqlparser.expression.LongValue; 10 import net.sf.jsqlparser.expression.LongValue;
@@ -22,24 +23,30 @@ import java.util.Set; @@ -22,24 +23,30 @@ import java.util.Set;
22 23
23 import static cn.iocoder.yudao.framework.common.util.collection.SetUtils.asSet; 24 import static cn.iocoder.yudao.framework.common.util.collection.SetUtils.asSet;
24 import static org.junit.jupiter.api.Assertions.assertEquals; 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 * 主要复用了 MyBatis Plus 的 TenantLineInnerInterceptorTest 的单元测试 31 * 主要复用了 MyBatis Plus 的 TenantLineInnerInterceptorTest 的单元测试
29 * 不过它的单元测试不是很规范，考虑到是复用的，所以暂时不进行修改~ 32 * 不过它的单元测试不是很规范，考虑到是复用的，所以暂时不进行修改~
30 * 33 *
31 * @author 芋道源码 34 * @author 芋道源码
32 */ 35 */
33 -public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest { 36 +public class DataPermissionRuleHandlerTest extends BaseMockitoUnitTest {
34 37
35 @InjectMocks 38 @InjectMocks
36 - private DataPermissionDatabaseInterceptor interceptor; 39 + private DataPermissionRuleHandler handler;
37 40
38 @Mock 41 @Mock
39 private DataPermissionRuleFactory ruleFactory; 42 private DataPermissionRuleFactory ruleFactory;
40 43
  44 + private DataPermissionInterceptor interceptor;
  45 +
41 @BeforeEach 46 @BeforeEach
42 public void setUp() { 47 public void setUp() {
  48 + interceptor = new DataPermissionInterceptor(handler);
  49 +
43 // 租户的数据权限规则 50 // 租户的数据权限规则
44 DataPermissionRule tenantRule = new DataPermissionRule() { 51 DataPermissionRule tenantRule = new DataPermissionRule() {
45 52
@@ -78,8 +85,8 @@ public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest @@ -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 @Test 92 @Test
@@ -448,7 +455,6 @@ public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest @@ -448,7 +455,6 @@ public class DataPermissionDatabaseInterceptorTest2 extends BaseMockitoUnitTest
448 assertEquals(targetSql, interceptor.parserSingle(sql, null)); 455 assertEquals(targetSql, interceptor.parserSingle(sql, null));
449 } 456 }
450 457
451 -  
452 // ========== 额外的测试 ========== 458 // ========== 额外的测试 ==========
453 459
454 @Test 460 @Test