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

package jnpf.database.plugins;

import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.extension.plugins.inner.TenantLineInnerInterceptor;
import jnpf.constant.MsgCode;
import jnpf.database.util.DynamicDataSourceUtil;
import jnpf.database.util.NotTenantPluginHolder;
import jnpf.exception.DataException;
import jnpf.util.TenantHolder;
import jodd.util.StringPool;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.*;
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.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

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

/**
 * Column模式租户插件
 * @author JNPF开发平台组
 * @user N
 * @copyright 引迈信息技术有限公司
 * @date 2022/10/14 10:29
 */
@Slf4j
public class MyTenantLineInnerInterceptor extends TenantLineInnerInterceptor implements ITenantPlugin {

    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
        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().isColumn()) {
            return;
        }
        if (NotTenantPluginHolder.isNotSwitch()) {
            NotTenantPluginHolder.clearNotSwitchFlag();
            return;
        }
        //非主库不切库
        if(!DynamicDataSourceUtil.isPrimaryDataSoure()){
            return;
        }
        //不绑定数据源的接口不切库
        /*if(NotTenantPluginHolder.isNotSwitchAlways()){
            return;
        }*/
        try {
            super.beforeQuery(executor, ms, parameter, rowBounds, resultHandler, boundSql);
        } catch (Exception e){
            //特殊语句解析失败
            if(log.isDebugEnabled()){
                log.debug("语句解析失败", e);
            }
        }
    }

    @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().isColumn()) {
            return;
        }
        if (NotTenantPluginHolder.isNotSwitch()) {
            NotTenantPluginHolder.clearNotSwitchFlag();
            return;
        }
        //非主库不切库
        if(!DynamicDataSourceUtil.isPrimaryDataSoure()){
            return;
        }
        //不绑定数据源的接口不切库
        /*if(NotTenantPluginHolder.isNotSwitchAlways()){
            return;
        }*/
        try {
            super.beforePrepare(sh, connection, transactionTimeout);
        } catch (Exception e){
            //特殊语句解析失败
            if(log.isDebugEnabled()){
                log.debug("语句解析失败", e);
            }
        }
    }

    @Override
    protected Column getAliasColumn(Table table) {
        return getAliasColumnWithFromItem(table);
    }

    protected Column getAliasColumnWithFromItem(FromItem table) {
        StringBuilder column = new StringBuilder();
        if (table.getAlias() != null) {
            column.append(table.getAlias().getName());
        }else{
            if(table instanceof Table){
                column.append(((Table)table).getName());
            }
        }
        column.append(StringPool.DOT).append(super.getTenantLineHandler().getTenantIdColumn());
        return new Column(column.toString());
    }

    protected void appendSelectItem(List<SelectItem<?>> selectItems, FromItem from) {
        if (CollectionUtils.isEmpty(selectItems)) {
            return;
        }
        if (selectItems.size() == 1) {
            SelectItem item = selectItems.get(0);
            Expression expression = item.getExpression();
            if (expression instanceof AllColumns) {
                return;
            }
        }
        // 改
        selectItems.add(new SelectItem<>(getAliasColumnWithFromItem(from)));
    }


    @Override
    protected void processInsertSelect(Select selectBody, final String whereSegment) {
        if (selectBody instanceof PlainSelect) {
            PlainSelect plainSelect = (PlainSelect)selectBody;
            FromItem fromItem = plainSelect.getFromItem();
            if (fromItem instanceof Table) {
                this.processPlainSelect(plainSelect, whereSegment);
                // 改
                this.appendSelectItem(plainSelect.getSelectItems(), fromItem);
            } else if (fromItem instanceof Select) {
                Select subSelect = (Select)fromItem;
                // 改
                this.appendSelectItem(plainSelect.getSelectItems(), fromItem);
                this.processInsertSelect(subSelect, whereSegment);
            }
        } else if (selectBody instanceof ParenthesedSelect) {
            ParenthesedSelect parenthesedSelect = (ParenthesedSelect)selectBody;
            this.processInsertSelect(parenthesedSelect.getSelect(), whereSegment);
        }

    }
}