更新 SaaSTenantSqlParser.java.ftl

SLN/%PUBPRJ%-core/src/main/java/%SYS_PKGPATH%/core/util/config/SaaSTenantSqlParser.java.ftl

@@ -6,12 +6,13 @@ package ${pub.getPKGCodeName()}.core.util.config;
 
 import com.baomidou.mybatisplus.core.toolkit.StringPool;
 import com.baomidou.mybatisplus.extension.plugins.tenant.TenantSqlParser;
+import net.sf.jsqlparser.expression.Expression;
+import net.sf.jsqlparser.expression.Function;
+import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
+import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
 import net.sf.jsqlparser.schema.Column;
 import net.sf.jsqlparser.schema.Table;
-import net.sf.jsqlparser.statement.select.FromItem;
-import net.sf.jsqlparser.statement.select.Join;
-import net.sf.jsqlparser.statement.select.PlainSelect;
-import net.sf.jsqlparser.statement.select.SelectExpressionItem;
+import net.sf.jsqlparser.statement.select.*;
 
 import java.util.List;
 
@@ -24,6 +25,7 @@ public class SaaSTenantSqlParser extends TenantSqlParser {
      * @param addColumn   是否添加租户列,insert into select语句中需要
      */
     protected void processPlainSelect(PlainSelect plainSelect, boolean addColumn) {
+        processSelectItem(plainSelect);
         FromItem fromItem = plainSelect.getFromItem();
         if (fromItem instanceof Table) {
             Table fromTable = (Table) fromItem;
@@ -46,6 +48,49 @@ public class SaaSTenantSqlParser extends TenantSqlParser {
                 processFromItem(j.getRightItem());
             });
         }
+        processWhere(plainSelect.getWhere());
     }
+
+    /**
+     * where 条件中包含select  添加租户id
+     */
+    protected void processWhere(Expression where) {
+        if (where == null)
+            return;
+        if (where instanceof AndExpression) {
+            processWhere(((AndExpression) where).getLeftExpression());
+            processWhere(((AndExpression) where).getRightExpression());
+        } else if (where instanceof OrExpression) {
+            processWhere(((OrExpression) where).getLeftExpression());
+            processWhere(((OrExpression) where).getRightExpression());
+        } else if (where instanceof Function) {
+            for (Expression e : ((Function) where).getParameters().getExpressions()) {
+                if (e instanceof SubSelect) {
+                    this.processSelectBody(((SubSelect) e).getSelectBody());
+                }
+            }
+        }
+    }
+
+    /**
+     * select 中包含 select  添加租户id
+     */
+    protected void processSelectItem(PlainSelect plainSelect) {
+        List<SelectItem> selectItems = plainSelect.getSelectItems();
+        for (SelectItem selectItem : selectItems) {
+            if (selectItem instanceof SelectExpressionItem) {
+                if(((SelectExpressionItem) selectItem).getExpression() instanceof SubSelect){
+                    this.processSelectBody(((SubSelect) ((SelectExpressionItem) selectItem).getExpression()).getSelectBody());
+                }
+            } else if (selectItem instanceof Function) {
+                for (Expression e : ((Function) selectItem).getParameters().getExpressions()) {
+                    if (e instanceof Select) {
+                        this.processSelectBody(((Select) e).getSelectBody());
+                    }
+                }
+            }
+        }
+    }
+
 }
 </#if>
\ No newline at end of file