package com.baomidou.mybatisplus.samples.tenant.config;

import org.mybatis.spring.annotation.MapperScan;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.baomidou.mybatisplus.autoconfigure.ConfigurationCustomizer;
import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
import com.baomidou.mybatisplus.extension.plugins.handler.TenantLineHandler;
import com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor;

import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;

/**
 * @author miemie
 * @since 2018-08-10
 */
@Configuration
@MapperScan("com.baomidou.mybatisplus.samples.tenant.mapper")
public class MybatisPlusConfig {

    /**
     * 新多租戶插件配置,一緩和二緩遵循mybatis的規(guī)則,需要設(shè)置 MybatisConfiguration#useDeprecatedExecutor = false 避免緩存萬(wàn)一出現(xiàn)問(wèn)題
     */
    @Bean
    public MybatisPlusInterceptor mybatisPlusInterceptor() {
        MybatisPlusInterceptor interceptor = new MybatisPlusInterceptor();
        interceptor.addInnerInterceptor(new TenantLineInnerInterceptor(new TenantLineHandler() {
            @Override
            public Expression getTenantId() {
                // 這里可以寫(xiě)自己系統(tǒng)的獲取租戶id的方法   比如下面自定義方法
                // return new LongValue(CurrentUserUtils.getTenantId());
                return new LongValue(1);
            }
            
            /**
             * 獲取租戶字段名
             * <p>
             * 默認(rèn)字段名叫: tenant_id
             *
             * @return 租戶字段名
             */
             @Override
            default String getTenantIdColumn() {
                // 如果該字段你不是固定的，請(qǐng)使用 SqlInjectionUtils.check 檢查安全性
                return "tenant_id";
            }
        
             /**
              * 根據(jù)表名判斷是否忽略拼接多租戶條件
              * <p>
              * 默認(rèn)都要進(jìn)行解析并拼接多租戶條件
              *
              * @param tableName 表名
              * @return 是否忽略, true:表示忽略，false:需要解析并拼接多租戶條件
              */
            @Override
            public boolean ignoreTable(String tableName) {
               String[] arr = new String[]{
                       "susCode",
                       "suName",
                       "suSex",
                       "suAge"
               };
               return ArrayUtil.contains(arr, tableName);
            }
        }));
        // 如果用了分頁(yè)插件注意先 add TenantLineInnerInterceptor 再 add PaginationInnerInterceptor
        // 用了分頁(yè)插件必須設(shè)置 MybatisConfiguration#useDeprecatedExecutor = false
//        interceptor.addInnerInterceptor(new PaginationInnerInterceptor());
        return interceptor;
    }

//    @Bean
//    public ConfigurationCustomizer configurationCustomizer() {
//        return configuration -> configuration.setUseDeprecatedExecutor(false);
//    }
}

package com.panpass.rebate.service.config;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.ReflectUtil;
import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.extension.handlers.AbstractSqlParserHandler;
import com.panpass.rebate.service.utils.CurrentRebateUserUtil;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.*;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.BaseStatementHandler;
import org.apache.ibatis.executor.statement.PreparedStatementHandler;
import org.apache.ibatis.executor.statement.RoutingStatementHandler;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.junit.jupiter.api.Order;
import org.springframework.core.Ordered;
import org.springframework.stereotype.Component;

import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.update.Update;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;



/**
 * @Description:
 * @Author: potato
 * @Date: 2024/4/8  9:49
 */


@Slf4j
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
//@Component
@Order(Ordered.HIGHEST_PRECEDENCE)
public class DataScopeInterceptor  extends AbstractSqlParserHandler implements Interceptor {

    /**
     * 是否對(duì)xml中SQL 進(jìn)行多租戶增強(qiáng)
     *  false   不進(jìn)行SQL增強(qiáng)
     *  true    進(jìn)行SQL增強(qiáng)
     *  配置 可轉(zhuǎn)移到nacos 中
     */
    private boolean openIngore = true;

    /**
     * 可以指定某些特定的XML中的SQL的方法進(jìn)行手寫(xiě)，不通過(guò)該攔截器增強(qiáng)
     */
    private List<String> ingoreList;


    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Object target = invocation.getTarget();
        //確保只有攔截的目標(biāo)對(duì)象是 StatementHandler 類(lèi)型時(shí)才執(zhí)行特定邏輯
        boolean flag = true;
        if (target instanceof StatementHandler) {
            StatementHandler statementHandler = (StatementHandler) target;
            StatementHandler delegate = (StatementHandler) ReflectUtil.getFieldValue(statementHandler, "delegate");

            // true 則需要SQL多租戶增強(qiáng)
            try {
                flag = isIngoreXML(invocation);
            } catch (Exception e) {
                log.info("StatementHandler 解析出現(xiàn)問(wèn)題，請(qǐng)注意查看"+e.getMessage());
//                e.printStackTrace();
            }
            if (flag) {
                // 獲取 BoundSql 對(duì)象，包含原始 SQL 語(yǔ)句
                BoundSql boundSql = statementHandler.getBoundSql();
                String originalSql = boundSql.getSql();
                String newSql = setEnvToStatement(originalSql);
                // 使用MetaObject對(duì)象將新的SQL語(yǔ)句設(shè)置到BoundSql對(duì)象中
                MetaObject metaObject = SystemMetaObject.forObject(boundSql);
                metaObject.setValue("sql", newSql);
            }

        }
        // 執(zhí)行SQL
        return invocation.proceed();
    }

    /**
     *
     *  根據(jù)xml名稱(chēng)判斷是否需要進(jìn)行xml 中SQL解析   可以拓展
     *
     * @param invocation
     * @return  fasle = 不進(jìn)行SQL增強(qiáng)   true 進(jìn)行SQL增強(qiáng)
     */
    public boolean isIngoreXML(Invocation invocation){
        // 不開(kāi)啟增強(qiáng)直接返回false
        if(!openIngore){
            return  false;
        }
        // 存在多層代理，SystemMetaObject獲取不確定，  借鑒PaginationInterceptor 中獲取Target的方法來(lái)獲取
        StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
        // SQL 解析
//        this.sqlParser(metaObject);
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");

//         ingoreList  為空，沒(méi)有設(shè)置需要忽略增強(qiáng)的方法名   模擬存在配置了忽略的方法
//         ingoreList =new ArrayList<>();
//         ingoreList.add("RebateFlowBulkMapper.flowList");
        if(CollectionUtil.isEmpty(ingoreList)){
            return true;
        }
        // Resource 示例 :   file [D:\bjzx\project\rebate\rebate-dao\target\classes\mapper\plan\PlanConfigMapper.xml]   xml路徑
        // id       示例 :   com.panpass.rebate.plan.persistent.mapper.PlanConfigMapper.getPageList
        // 確保Resource 來(lái)自xml,并且配置的存在忽略的方法名例如： PlanConfigMapper.getPageList，返回false  不進(jìn)行增強(qiáng)處理;
        for (String menthodName : ingoreList) {
            if (mappedStatement.getResource().contains("xml") && mappedStatement.getId().contains(menthodName)) {
                return false;
            }
        }

        return true;
    }

    public static void printFields(Object obj) {
        Class clazz = obj.getClass();
        Field[] fields = clazz.getDeclaredFields();
        for (Field field : fields) {
            field.setAccessible(true);
            String fieldName = field.getName();
            Object fieldValue = null;
            try {
                fieldValue = field.get(obj);
            } catch (IllegalAccessException e) {
                e.printStackTrace();
            }
            System.out.println(fieldName + ": " + fieldValue);
        }
    }

    private String setEnvToStatement(String originalSql) {
        net.sf.jsqlparser.statement.Statement statement;
        try {
            statement = CCJSqlParserUtil.parse(originalSql);
        } catch (JSQLParserException e) {
            throw new RuntimeException("EnvironmentVariableInterceptor::SQL語(yǔ)句解析異常："+originalSql);
        }

        if (statement instanceof Select) {
            //TODO需要遞歸處理
            Select select = (Select) statement;
            PlainSelect selectBody = (PlainSelect) select.getSelectBody();
            // 遞歸處理子查詢(xún)
            processSubQuery(selectBody);

            return select.toString();
        }
//        else if (statement instanceof Insert) {
//            Insert insert = (Insert) statement;
//            setEnvToInsert(insert);
//
//            return insert.toString();
//        } else if (statement instanceof Update) {
//            Update update = (Update) statement;
//            Expression newWhereExpression = setEnvToWhereExpression(update.getWhere(),null);
//            update.setWhere(newWhereExpression);
//
//            return update.toString();
//        } else if (statement instanceof Delete) {
//            Delete delete = (Delete) statement;
//            Expression newWhereExpression = setEnvToWhereExpression(delete.getWhere(),null);
//            delete.setWhere(newWhereExpression);
//
//            return delete.toString();
//        }
        return originalSql;
    }

    private  void processSubQuery(PlainSelect selectBody ) {

            if(selectBody == null){
                return;
            }
        if (selectBody.getFromItem() instanceof Table) {
            Expression newWhereExpression;
            if (selectBody.getJoins() == null || selectBody.getJoins().isEmpty()) {
                newWhereExpression = setEnvToWhereExpression(selectBody.getWhere(), null);
            } else {
                newWhereExpression = multipleTableJoinWhereExpression(selectBody);
            }
            selectBody.setWhere(newWhereExpression);
        } else {
            // 處理子查詢(xún)
            SelectBody subSelectBody = ((SubSelect) selectBody.getFromItem()).getSelectBody();
            processSubQuery((PlainSelect) subSelectBody);
//            selectBody.setFromItem((FromItem) processSubQuery((PlainSelect) subSelectBody));
        }
//        return selectBody;
    }

    //非遞查詢(xún)
    private  String setEnvToStatement2(String originalSql) {
        net.sf.jsqlparser.statement.Statement statement;
        try {
            statement = CCJSqlParserUtil.parse(originalSql);
        } catch (JSQLParserException e) {
            throw new RuntimeException("EnvironmentVariableInterceptor::SQL語(yǔ)句解析異常："+originalSql);
        }
        if (statement instanceof Select) {
            Select select = (Select) statement;
            PlainSelect selectBody = (PlainSelect) select.getSelectBody();
            if (selectBody.getFromItem() instanceof Table) {
                Expression newWhereExpression;
                if (selectBody.getJoins() == null || selectBody.getJoins().isEmpty()) {
                    newWhereExpression = setEnvToWhereExpression(selectBody.getWhere(), null);
                } else {
                    // 如果是多表關(guān)聯(lián)查詢(xún)，在關(guān)聯(lián)查詢(xún)中新增每個(gè)表的環(huán)境變量條件
                    newWhereExpression = multipleTableJoinWhereExpression(selectBody);
                }
                // 將新的where設(shè)置到Select中
                selectBody.setWhere(newWhereExpression);
            } else if (selectBody.getFromItem() instanceof SubSelect) {
                // 如果是子查詢(xún)，在子查詢(xún)中新增環(huán)境變量條件
                // 當(dāng)前方法只能處理單層子查詢(xún)，如果有多層級(jí)的子查詢(xún)的場(chǎng)景需要通過(guò)遞歸設(shè)置環(huán)境變量
                SubSelect subSelect = (SubSelect) selectBody.getFromItem();
                PlainSelect subSelectBody = (PlainSelect) subSelect.getSelectBody();
                Expression newWhereExpression = setEnvToWhereExpression(subSelectBody.getWhere(), null);
                subSelectBody.setWhere(newWhereExpression);
            }

            // 獲得修改后的語(yǔ)句
            return select.toString();
        } else if (statement instanceof Insert) {
            Insert insert = (Insert) statement;
            setEnvToInsert(insert);

            return insert.toString();
        } else if (statement instanceof Update) {
            Update update = (Update) statement;
            Expression newWhereExpression = setEnvToWhereExpression(update.getWhere(),null);
            // 將新的where設(shè)置到Update中
            update.setWhere(newWhereExpression);

            return update.toString();
        } else if (statement instanceof Delete) {
            Delete delete = (Delete) statement;
            Expression newWhereExpression = setEnvToWhereExpression(delete.getWhere(),null);
            // 將新的where設(shè)置到delete中
            delete.setWhere(newWhereExpression);

            return delete.toString();
        }
        return originalSql;
    }

    /**
     * 將需要隔離的字段加入到SQL的Where語(yǔ)法樹(shù)中
     * @param whereExpression SQL的Where語(yǔ)法樹(shù)
     * @param alias 表別名
     * @return 新的SQL Where語(yǔ)法樹(shù)
     */
    private Expression setEnvToWhereExpression(Expression whereExpression, String alias) {
        // 添加SQL語(yǔ)法樹(shù)的一個(gè)where分支，并添加環(huán)境變量條件

        EqualsTo envEquals = new EqualsTo();
        envEquals.setLeftExpression(new Column(StringUtils.isNotBlank(alias) ? String.format("%s.tenant_id", alias) : "tenant_id"));
        envEquals.setRightExpression(new LongValue(CurrentRebateUserUtil.getTenantId() == null ? 1 :CurrentRebateUserUtil.getTenantId()));
        if (whereExpression == null){
            return envEquals;
        } else {
            AndExpression andExpression = new AndExpression(whereExpression,envEquals);
            // 將新的where條件加入到原where條件的右分支樹(shù)
            andExpression.setRightExpression(envEquals);
            andExpression.setLeftExpression(whereExpression);
            return andExpression;
        }
    }

    /**
     * 多表關(guān)聯(lián)查詢(xún)時(shí)，給關(guān)聯(lián)的所有表加入環(huán)境隔離條件
     * @param selectBody select語(yǔ)法樹(shù)
     * @return 新的SQL Where語(yǔ)法樹(shù)
     */
    private Expression multipleTableJoinWhereExpression(PlainSelect selectBody){
        Table mainTable = (Table) selectBody.getFromItem();
        String mainTableAlias = mainTable.getAlias().getName();
        // 將 t1.tenant_id = tenant_id 的條件添加到where中
        Expression newWhereExpression = setEnvToWhereExpression(selectBody.getWhere(), mainTableAlias);
        List<Join> joins = selectBody.getJoins();
        for (Join join : joins) {
            FromItem joinRightItem = join.getRightItem();
            if (joinRightItem instanceof Table) {
                Table joinTable = (Table) joinRightItem;
                String joinTableAlias = joinTable.getAlias().getName();
                // 將每一個(gè)join的 tx.env = ENV 的條件添加到where中
                newWhereExpression = setEnvToWhereExpression(newWhereExpression, joinTableAlias);
            }
        }
        return newWhereExpression;
    }

    /**
     * 新增數(shù)據(jù)時(shí)，插入tenant_id字段
     * @param insert Insert 語(yǔ)法樹(shù)
     */
    private void setEnvToInsert(Insert insert) {
        // 添加tenant_id列
        List<Column> columns = insert.getColumns();
        for (Column column : columns) {
            //若存在，不進(jìn)行處理
            if (column.getColumnName().equals("tenant_id")) {
                return;
            }
        }
        columns.add(new Column("tenant_id"));
        // values中添加環(huán)境變量值

        // 獲取插入值列表
        ItemsList itemsList = insert.getItemsList();
        if(itemsList instanceof MultiExpressionList ){
            List<ExpressionList> exprList = ((MultiExpressionList) itemsList).getExprList();
            for (ExpressionList expressionList : exprList) {
                expressionList.getExpressions().add(new LongValue(CurrentRebateUserUtil.getTenantId() == null ? 1 :CurrentRebateUserUtil.getTenantId()));
            }
        } else if (itemsList instanceof SubSelect) {
            // 處理子查詢(xún)
            log.info("子查詢(xún)插入語(yǔ)句業(yè)務(wù)");
        }

    }

    //測(cè)試邏輯使用的
    public static void main(String[] args) throws JSQLParserException {

        String slectSql = "select * from (select id from (select 1 from table3 where id =2))";
//        String s = setEnvToStatement(slectSql);

        String sql = "INSERT INTO my_table (column1, column2) VALUES (value1, value2), (value3, value4)";

        try {
            // 解析 INSERT 語(yǔ)句
            Statement statement = CCJSqlParserUtil.parse(sql);

            // 判斷是否是 INSERT 語(yǔ)句
            if (statement instanceof Insert) {
                Insert insert = (Insert) statement;
                List<Column> columns = insert.getColumns();

                columns.add(new Column("tenant_id"));
                // 獲取插入值列表
                ItemsList itemsList = insert.getItemsList();
                if(itemsList instanceof MultiExpressionList ){
                    itemsList =(MultiExpressionList)itemsList;
                    List<ExpressionList> exprList = ((MultiExpressionList) itemsList).getExprList();
                    for (ExpressionList expressionList : exprList) {
                        expressionList.getExpressions().add(new LongValue(CurrentRebateUserUtil.getTenantId() == null ? 1 :CurrentRebateUserUtil.getTenantId()));
                    }
                } else if (itemsList instanceof SubSelect) {
                    // 處理子查詢(xún)
                    log.info("子查詢(xún)插入語(yǔ)句使用");
                }
            }

        } catch (JSQLParserException e) {
            e.printStackTrace();
        }
    }




    /**
     * 生成攔截對(duì)象的代理
     *
     * @param target 目標(biāo)對(duì)象
     * @return 代理對(duì)象
     */
    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    /**
     * mybatis配置的屬性
     *
     * @param properties mybatis配置的屬性
     */
    @Override
    public void setProperties(Properties properties) {

    }

}

package com.panpass.rebate.service.config;


import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.session.SqlSessionFactory;

import org.springframework.beans.factory.annotation.Autowired;

import org.springframework.context.annotation.Configuration;

import javax.annotation.PostConstruct;
import java.util.List;

/**
 * @Description:
 * @Author: potato
 * @Date: 2024/4/3  13:49
 */

@Configuration
public class MyBatisConfig {





    @Autowired
    private List<SqlSessionFactory> sqlSessionFactoryList;

    // 只執(zhí)行一次
    @PostConstruct
    public void addDefaultTimeInterceptor() {
        /**
         * Mybatis攔截器可以使用@Component注解也可以在這里進(jìn)行配置
         * 在這里配置可以控制攔截器的執(zhí)行順序，所以注意去掉@Component注解
         */
        for (SqlSessionFactory sqlSessionFactory : sqlSessionFactoryList) {
            org.apache.ibatis.session.Configuration configuration = sqlSessionFactory.getConfiguration();
            List<Interceptor> interceptors = configuration.getInterceptors();
            // 最后添加的會(huì)更早執(zhí)行
            
            configuration.addInterceptor(new PaginationInterceptor());
            configuration.addInterceptor(new DataScopeInterceptor());
        }
    }

}