实现 DataPermissionRuleFactoryImpl，并增加相关单测

YunaiV
1 parent f9b15fe7
Showing 6 changed files with 284 additions and 15 deletions
yudao-framework/yudao-spring-boot-starter-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/core/aop/DataPermissionAnnotationInterceptor.java
yudao-framework/yudao-spring-boot-starter-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/core/aop/DataPermissionContextHolder.java
yudao-framework/yudao-spring-boot-starter-data-permission/src/main/java/cn/iocoder/yudao/framework/datapermission/core/rule/DataPermissionRuleFactoryImpl.java
yudao-framework/yudao-spring-boot-starter-data-permission/src/test/java/cn/iocoder/yudao/framework/datapermission/core/aop/DataPermissionContextHolderTest.java
yudao-framework/yudao-spring-boot-starter-data-permission/src/test/java/cn/iocoder/yudao/framework/datapermission/core/rule/DataPermissionRuleFactoryImplTest.java
@@ -34,7 +34,7 @@ public class DataPermissionAnnotationInterceptor implements MethodInterceptor {
         // 入栈
         DataPermission dataPermission = this.findAnnotation(methodInvocation);
         if (dataPermission != null) {
-            DataPermissionContextHolder.push(dataPermission);
+            DataPermissionContextHolder.add(dataPermission);
         }
         try {
             // 执行逻辑
@@ -42,7 +42,7 @@ public class DataPermissionAnnotationInterceptor implements MethodInterceptor {
         } finally {
             // 出栈
             if (dataPermission != null) {
-                DataPermissionContextHolder.poll();
+                DataPermissionContextHolder.remove();
             }
         }
     }
@@ -3,8 +3,8 @@ package cn.iocoder.yudao.framework.datapermission.core.aop;
 import cn.iocoder.yudao.framework.datapermission.core.annotation.DataPermission;
 import com.alibaba.ttl.TransmittableThreadLocal;
  
-import java.util.ArrayDeque;
-import java.util.Deque;
+import java.util.LinkedList;
+import java.util.List;
  
 /**
  * {@link DataPermission} 注解的 Context 上下文
@@ -13,16 +13,19 @@ import java.util.Deque;
  */
 public class DataPermissionContextHolder {
  
-    private static final ThreadLocal<Deque<DataPermission>> DATA_PERMISSIONS =
-            TransmittableThreadLocal.withInitial(ArrayDeque::new);
+    /**
+     * 使用 List 的原因，可能存在方法的嵌套调用
+     */
+    private static final ThreadLocal<LinkedList<DataPermission>> DATA_PERMISSIONS =
+            TransmittableThreadLocal.withInitial(LinkedList::new);
  
     /**
      * 获得当前的 DataPermission 注解
      *
      * @return DataPermission 注解
      */
-    public static DataPermission peek() {
-        return DATA_PERMISSIONS.get().remove();
+    public static DataPermission get() {
+        return DATA_PERMISSIONS.get().peekLast();
     }
  
     /**
@@ -30,8 +33,8 @@ public class DataPermissionContextHolder {
      *
      * @param dataPermission DataPermission 注解
      */
-    public static void push(DataPermission dataPermission) {
-        DATA_PERMISSIONS.get().push(dataPermission);
+    public static void add(DataPermission dataPermission) {
+        DATA_PERMISSIONS.get().addLast(dataPermission);
     }
  
     /**
@@ -39,13 +42,31 @@ public class DataPermissionContextHolder {
      *
      * @return DataPermission 注解
      */
-    public static DataPermission poll() {
-        DataPermission dataPermission = DATA_PERMISSIONS.get().poll();
+    public static DataPermission remove() {
+        DataPermission dataPermission = DATA_PERMISSIONS.get().removeLast();
         // 无元素时，清空 ThreadLocal
-        if (dataPermission == null) {
+        if (DATA_PERMISSIONS.get().isEmpty()) {
             DATA_PERMISSIONS.remove();
         }
         return dataPermission;
     }
  
+    /**
+     * 获得所有 DataPermission
+     *
+     * @return DataPermission 队列
+     */
+    public static List<DataPermission> getAll() {
+        return DATA_PERMISSIONS.get();
+    }
+
+    /**
+     * 清空上下文
+     *
+     * 目前仅仅用于单测
+     */
+    public static void clear() {
+        DATA_PERMISSIONS.remove();
+    }
+
 }
 package cn.iocoder.yudao.framework.datapermission.core.rule;
  
+import cn.hutool.core.collection.CollUtil;
+import cn.hutool.core.util.ArrayUtil;
+import cn.iocoder.yudao.framework.datapermission.core.annotation.DataPermission;
+import cn.iocoder.yudao.framework.datapermission.core.aop.DataPermissionContextHolder;
 import lombok.RequiredArgsConstructor;
  
+import java.util.Collections;
 import java.util.List;
+import java.util.stream.Collectors;
  
+/**
+ * 默认的 DataPermissionRuleFactoryImpl 实现类
+ * 支持通过 {@link DataPermissionContextHolder} 过滤数据权限
+ *
+ * @author 芋道源码
+ */
 @RequiredArgsConstructor
 public class DataPermissionRuleFactoryImpl implements DataPermissionRuleFactory {
  
@@ -17,9 +29,34 @@ public class DataPermissionRuleFactoryImpl implements DataPermissionRuleFactory 
         return rules;
     }
  
-    @Override
+    @Override // mappedStatementId 参数，暂时没有用。以后，可以基于 mappedStatementId + DataPermission 进行缓存
     public List<DataPermissionRule> getDataPermissionRule(String mappedStatementId) {
-        return null;
+        // 1. 无数据权限
+        if (CollUtil.isEmpty(rules)) {
+            return Collections.emptyList();
+        }
+        // 2. 未配置，则默认开启
+        DataPermission dataPermission = DataPermissionContextHolder.get();
+        if (dataPermission == null) {
+            return rules;
+        }
+        // 3. 已配置，但禁用
+        if (!dataPermission.enable()) {
+            return Collections.emptyList();
+        }
+
+        // 4. 已配置，只选择部分规则
+        if (ArrayUtil.isNotEmpty(dataPermission.includeRules())) {
+            return rules.stream().filter(rule -> ArrayUtil.contains(dataPermission.includeRules(), rule.getClass()))
+                    .collect(Collectors.toList()); // 一般规则不会太多，所以不采用 HashSet 查询
+        }
+        // 5. 已配置，只排除部分规则
+        if (ArrayUtil.isNotEmpty(dataPermission.excludeRules())) {
+            return rules.stream().filter(rule -> !ArrayUtil.contains(dataPermission.excludeRules(), rule.getClass()))
+                    .collect(Collectors.toList()); // 一般规则不会太多，所以不采用 HashSet 查询
+        }
+        // 6. 已配置，全部规则
+        return rules;
     }
  
 }
+package cn.iocoder.yudao.framework.datapermission.core.aop;
+
+import cn.iocoder.yudao.framework.datapermission.core.annotation.DataPermission;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.mockito.Mockito.mock;
+
+/**
+ * {@link DataPermissionContextHolder} 的单元测试
+ *
+ * @author 芋道源码
+ */
+class DataPermissionContextHolderTest {
+
+    @BeforeEach
+    public void setUp() {
+        DataPermissionContextHolder.clear();
+    }
+
+    @Test
+    public void testGet() {
+        // mock 方法
+        DataPermission dataPermission01 = mock(DataPermission.class);
+        DataPermissionContextHolder.add(dataPermission01);
+        DataPermission dataPermission02 = mock(DataPermission.class);
+        DataPermissionContextHolder.add(dataPermission02);
+
+        // 调用
+        DataPermission result = DataPermissionContextHolder.get();
+        // 断言
+        assertSame(result, dataPermission02);
+    }
+
+    @Test
+    public void testPush() {
+        // 调用
+        DataPermission dataPermission01 = mock(DataPermission.class);
+        DataPermissionContextHolder.add(dataPermission01);
+        DataPermission dataPermission02 = mock(DataPermission.class);
+        DataPermissionContextHolder.add(dataPermission02);
+        // 断言
+        DataPermission first = DataPermissionContextHolder.getAll().get(0);
+        DataPermission second = DataPermissionContextHolder.getAll().get(1);
+        assertSame(dataPermission01, first);
+        assertSame(dataPermission02, second);
+    }
+
+    @Test
+    public void testRemove() {
+        // mock 方法
+        DataPermission dataPermission01 = mock(DataPermission.class);
+        DataPermissionContextHolder.add(dataPermission01);
+        DataPermission dataPermission02 = mock(DataPermission.class);
+        DataPermissionContextHolder.add(dataPermission02);
+
+        // 调用
+        DataPermission result = DataPermissionContextHolder.remove();
+        // 断言
+        assertSame(result, dataPermission02);
+        assertEquals(1, DataPermissionContextHolder.getAll().size());
+    }
+
+}
+package cn.iocoder.yudao.framework.datapermission.core.rule;
+
+import cn.iocoder.yudao.framework.datapermission.core.annotation.DataPermission;
+import cn.iocoder.yudao.framework.datapermission.core.aop.DataPermissionContextHolder;
+import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
+import net.sf.jsqlparser.expression.Alias;
+import net.sf.jsqlparser.expression.Expression;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.InjectMocks;
+import org.mockito.Spy;
+import org.springframework.core.annotation.AnnotationUtils;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Set;
+
+import static cn.iocoder.yudao.framework.test.core.util.RandomUtils.randomString;
+import static org.junit.jupiter.api.Assertions.*;
+
+/**
+ * {@link DataPermissionRuleFactoryImpl} 单元测试
+ *
+ * @author 芋道源码
+ */
+class DataPermissionRuleFactoryImplTest extends BaseMockitoUnitTest {
+
+    @InjectMocks
+    private DataPermissionRuleFactoryImpl dataPermissionRuleFactory;
+
+    @Spy
+    private List<DataPermissionRule> rules = Arrays.asList(new DataPermissionRule01(),
+            new DataPermissionRule02());
+
+    @BeforeEach
+    public void setUp() {
+        DataPermissionContextHolder.clear();
+    }
+
+    @Test
+    public void testGetDataPermissionRule_02() {
+        // 准备参数
+        String mappedStatementId = randomString();
+
+        // 调用
+        List<DataPermissionRule> result = dataPermissionRuleFactory.getDataPermissionRule(mappedStatementId);
+        // 断言
+        assertSame(rules, result);
+    }
+
+    @Test
+    public void testGetDataPermissionRule_03() {
+        // 准备参数
+        String mappedStatementId = randomString();
+        // mock 方法
+        DataPermissionContextHolder.add(AnnotationUtils.findAnnotation(TestClass03.class, DataPermission.class));
+
+        // 调用
+        List<DataPermissionRule> result = dataPermissionRuleFactory.getDataPermissionRule(mappedStatementId);
+        // 断言
+        assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void testGetDataPermissionRule_04() {
+        // 准备参数
+        String mappedStatementId = randomString();
+        // mock 方法
+        DataPermissionContextHolder.add(AnnotationUtils.findAnnotation(TestClass04.class, DataPermission.class));
+
+        // 调用
+        List<DataPermissionRule> result = dataPermissionRuleFactory.getDataPermissionRule(mappedStatementId);
+        // 断言
+        assertEquals(1, result.size());
+        assertEquals(DataPermissionRule01.class, result.get(0).getClass());
+    }
+
+    @Test
+    public void testGetDataPermissionRule_05() {
+        // 准备参数
+        String mappedStatementId = randomString();
+        // mock 方法
+        DataPermissionContextHolder.add(AnnotationUtils.findAnnotation(TestClass05.class, DataPermission.class));
+
+        // 调用
+        List<DataPermissionRule> result = dataPermissionRuleFactory.getDataPermissionRule(mappedStatementId);
+        // 断言
+        assertEquals(1, result.size());
+        assertEquals(DataPermissionRule02.class, result.get(0).getClass());
+    }
+
+    @Test
+    public void testGetDataPermissionRule_06() {
+        // 准备参数
+        String mappedStatementId = randomString();
+        // mock 方法
+        DataPermissionContextHolder.add(AnnotationUtils.findAnnotation(TestClass06.class, DataPermission.class));
+
+        // 调用
+        List<DataPermissionRule> result = dataPermissionRuleFactory.getDataPermissionRule(mappedStatementId);
+        // 断言
+        assertSame(rules, result);
+    }
+
+    @DataPermission(enable = false)
+    static class TestClass03 {}
+
+    @DataPermission(includeRules = DataPermissionRule01.class)
+    static class TestClass04 {}
+
+    @DataPermission(excludeRules = DataPermissionRule01.class)
+    static class TestClass05 {}
+
+    @DataPermission
+    static class TestClass06 {}
+
+    static class DataPermissionRule01 implements DataPermissionRule {
+
+        @Override
+        public Set<String> getTableNames() {
+            return null;
+        }
+
+        @Override
+        public Expression getExpression(String tableName, Alias tableAlias) {
+            return null;
+        }
+
+    }
+
+    static class DataPermissionRule02 implements DataPermissionRule {
+
+        @Override
+        public Set<String> getTableNames() {
+            return null;
+        }
+
+        @Override
+        public Expression getExpression(String tableName, Alias tableAlias) {
+            return null;
+        }
+
+    }
+
+}