jnpf-common-v6x/jnpf-boot-common/jnpf-common-database/target/classes/jnpf/database/plugins/MySchemaInnerInterceptor.java

package jnpf.database.plugins;

import com.baomidou.mybatisplus.core.plugins.InterceptorIgnoreHelper;
import com.baomidou.mybatisplus.core.toolkit.ExceptionUtils;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.baomidou.mybatisplus.core.toolkit.TableNameParser;
import com.baomidou.mybatisplus.extension.plugins.handler.TableNameHandler;
import com.baomidou.mybatisplus.extension.plugins.inner.DynamicTableNameInnerInterceptor;
import jnpf.constant.MsgCode;
import jnpf.database.util.DataSourceUtil;
import jnpf.database.util.DynamicDataSourceUtil;
import jnpf.database.util.NotTenantPluginHolder;
import jnpf.exception.DataException;
import jnpf.util.TenantHolder;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.beans.factory.annotation.Autowired;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * Schema模式租户插件
 * @author JNPF开发平台组
 * @user N
 * @copyright 引迈信息技术有限公司
 * @date 2022/10/14 10:35
 */
@Getter
@Setter
@NoArgsConstructor
@AllArgsConstructor
@Slf4j
public class MySchemaInnerInterceptor extends DynamicTableNameInnerInterceptor implements ITenantPlugin {

    private Runnable hook;

    public void setHook(Runnable hook) {
        this.hook = hook;
    }

    /**
     * 表名处理器，是否处理表名的情况都在该处理器中自行判断
     */
    private TableNameHandler tableNameHandler;
    @Autowired
    private DataSourceUtil dataSourceUtil;

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        if(TenantHolder.getLocalTenantCache() == null){
            printNoTenant(v -> log.warn("未设置租户信息, 禁止查询数据库, {}, {}, {}, {}", v.getUserId(), v.getUrl(), v.getToken(), v.getStack()));
            //未设置租户信息不允许操作数据库
            throw new DataException(MsgCode.LOG113.get());
        }
        //租户指定数据源不处理
        if (!TenantHolder.getLocalTenantCache().isSchema()) {
            return;
        }
        if (NotTenantPluginHolder.isNotSwitch()) {
            NotTenantPluginHolder.clearNotSwitchFlag();
            return;
        }
        //非主库不切库
        if(!DynamicDataSourceUtil.isPrimaryDataSoure()){
            return;
        }
        //不绑定数据源的接口不切库
        /*if(NotTenantPluginHolder.isNotSwitchAlways()){
            return;
        }*/
        PluginUtils.MPBoundSql mpBs = PluginUtils.mpBoundSql(boundSql);
        if (!InterceptorIgnoreHelper.willIgnoreDynamicTableName(ms.getId())) {
            // 非忽略执行
            mpBs.sql(this.changeTable(mpBs.sql()));
        }
    }

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        if(TenantHolder.getLocalTenantCache() == null){
            printNoTenant(v -> log.warn("未设置租户信息, 禁止查询数据库, {}, {}, {}, {}", v.getUserId(), v.getUrl(), v.getToken(), v.getStack()));
            //未设置租户信息不允许操作数据库
            throw new DataException(MsgCode.LOG113.get());
        }
        //租户指定数据源不处理
        if (!TenantHolder.getLocalTenantCache().isSchema()) {
            return;
        }
        if (NotTenantPluginHolder.isNotSwitch()) {
            NotTenantPluginHolder.clearNotSwitchFlag();
            return;
        }
        //非主库不切库
        if(!DynamicDataSourceUtil.isPrimaryDataSoure()){
            return;
        }
        //不绑定数据源的接口不切库
        /*if(NotTenantPluginHolder.isNotSwitchAlways()){
            return;
        }*/
        PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
        MappedStatement ms = mpSh.mappedStatement();
        SqlCommandType sct = ms.getSqlCommandType();
        if (sct == SqlCommandType.INSERT || sct == SqlCommandType.UPDATE || sct == SqlCommandType.DELETE) {
            if (!InterceptorIgnoreHelper.willIgnoreDynamicTableName(ms.getId())) {
                // 非忽略执行
                PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql();
                mpBs.sql(this.changeTable(mpBs.sql()));

                /*
                //只有修改、新增、删除数据才切库
                //先保留代码看后续需求
                SQLStatement sqlStatement = SQLParserUtils.createSQLStatementParser(mpBs.sql(), DbTypeUtil.getDb(dataSourceUtil).getDruidDbType()).parseStatement();
                if (sqlStatement instanceof SQLSelectStatement
                        || sqlStatement instanceof SQLUpdateStatement
                        || sqlStatement instanceof SQLInsertStatement
                        || sqlStatement instanceof SQLDeleteStatement) {
                    mpBs.sql(this.changeTable(mpBs.sql()));
                }*/
            }
        }
    }

    public String changeTable(String sql) {
        ExceptionUtils.throwMpe(null == tableNameHandler, "Please implement TableNameHandler processing logic");
        TableNameParser parser = new TableNameParser(sql);
        List<TableNameParser.SqlToken> names = new ArrayList<>();
        parser.accept(names::add);
        StringBuilder builder = new StringBuilder();
        int last = 0;
        for (TableNameParser.SqlToken name : names) {
            int start = name.getStart();
            if (start != last) {
                builder.append(sql, last, start);
                builder.append(tableNameHandler.dynamicTableName(sql, name.getValue()));
            }
            last = name.getEnd();
        }
        if (last != sql.length()) {
            builder.append(sql.substring(last));
        }
        if (hook != null) {
            hook.run();
        }
        return builder.toString();
    }
}