fixed github issues/3482 数据权限处理器支持 union all

mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/plugins/inner/DataPermissionInterceptor.java

@@ -23,6 +23,8 @@ import lombok.*;
 import net.sf.jsqlparser.expression.Expression;
 import net.sf.jsqlparser.statement.select.PlainSelect;
 import net.sf.jsqlparser.statement.select.Select;
+import net.sf.jsqlparser.statement.select.SelectBody;
+import net.sf.jsqlparser.statement.select.SetOperationList;
 import org.apache.ibatis.executor.Executor;
 import org.apache.ibatis.mapping.BoundSql;
 import org.apache.ibatis.mapping.MappedStatement;
@@ -30,6 +32,7 @@ import org.apache.ibatis.session.ResultHandler;
 import org.apache.ibatis.session.RowBounds;
 
 import java.sql.SQLException;
+import java.util.List;
 
 /**
  * 数据权限处理器
@@ -55,8 +58,24 @@ public class DataPermissionInterceptor extends JsqlParserSupport implements Inne
 
     @Override
     protected void processSelect(Select select, int index, String sql, Object obj) {
-        PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
-        Expression sqlSegment = dataPermissionHandler.getSqlSegment(plainSelect.getWhere(), (String) obj);
+        SelectBody selectBody = select.getSelectBody();
+        if (selectBody instanceof PlainSelect) {
+            this.setWhere((PlainSelect) selectBody, (String) obj);
+        } else if (selectBody instanceof SetOperationList) {
+            SetOperationList setOperationList = (SetOperationList) selectBody;
+            List<SelectBody> selectBodyList = setOperationList.getSelects();
+            selectBodyList.forEach(s -> this.setWhere((PlainSelect) s, (String) obj));
+        }
+    }
+
+    /**
+     * 设置 where 条件
+     *
+     * @param plainSelect  查询对象
+     * @param whereSegment 查询条件片段
+     */
+    protected void setWhere(PlainSelect plainSelect, String whereSegment) {
+        Expression sqlSegment = dataPermissionHandler.getSqlSegment(plainSelect.getWhere(), whereSegment);
         if (null != sqlSegment) {
             plainSelect.setWhere(sqlSegment);
         }

mybatis-plus-extension/src/test/java/com/baomidou/mybatisplus/extension/plugins/inner/DataPermissionInterceptorTest.java

@@ -80,6 +80,12 @@ public class DataPermissionInterceptorTest {
             "SELECT * FROM sys_role WHERE company_id = 6 AND companyId IN (1, 2, 3)");
     }
 
+    @Test
+    void test3unionAll() {
+        assertSql(TEST_3, "select * from sys_role where company_id=6 union all select * from sys_role where company_id=7",
+            "SELECT * FROM sys_role WHERE company_id = 6 AND companyId IN (1, 2, 3) UNION ALL SELECT * FROM sys_role WHERE company_id = 7 AND companyId IN (1, 2, 3)");
+    }
+
     @Test
     void test4() {
         assertSql(TEST_4, "select * from sys_role where id=3",