Mybatis plus 动态表名插件开发

import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import com.rookie.mybatis.study.entity.DynamicTableInfo;
import com.rookie.mybatis.study.entity.DynamicTableName;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.Join;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.update.Update;
import org.apache.commons.lang3.ObjectUtils;
import org.apache.commons.lang3.StringUtils;
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.plugin.*;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;

import java.sql.Connection;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.stream.Collectors;

/**
 * @author rookie
 */
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class DynamicTableInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
        MetaObject metaObject = SystemMetaObject.forObject(statementHandler);

        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
//        只处理增删改查语句
        if (SqlCommandType.SELECT != mappedStatement.getSqlCommandType()
                && SqlCommandType.INSERT != mappedStatement.getSqlCommandType()
                && SqlCommandType.DELETE != mappedStatement.getSqlCommandType()
                && SqlCommandType.UPDATE != mappedStatement.getSqlCommandType()) {
            return invocation.proceed();
        }

        // 针对定义了rowBounds，做为mapper接口方法的参数
        BoundSql boundSql = (BoundSql) metaObject.getValue("delegate.boundSql");
        Object paramObj = boundSql.getParameterObject();

        DynamicTableInfo info = extractInfo(paramObj);

        Map<String, String> tableMap = TablePluginContainer.TABLE_NAME_THREAD_LOCAL.get();

//        如果参数或者 ThreadLocal 中带有需要替换的表名
        boolean replaceProcess = (info == null || ObjectUtils.isEmpty(info.getTableNames())) && tableMap == null;

        if (replaceProcess) {
            return invocation.proceed();
        }

        Map<String, String> tableNameMap = null;

//        以 ThreadLocal 中的数据优先
        if (tableMap != null) {
            tableNameMap = tableMap;
        } else {
            tableNameMap = info.getTableNames().stream().collect(Collectors.toMap(DynamicTableName::getCurrentTable,
                    DynamicTableName::getItemTable, (v1, v2) -> v1));
        }

        String originalSql = boundSql.getSql();

        Statement statement = CCJSqlParserUtil.parse(originalSql);

//        替换相应表名
        originalSql = alterTableName(mappedStatement, statement, tableNameMap);

        metaObject.setValue("delegate.boundSql.sql", originalSql);
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof StatementHandler) {
            return Plugin.wrap(target, this);
        }
        return target;
    }

    @Override
    public void setProperties(Properties properties) {

    }

    private static Table alterTableName(Table table, Map<String, String> tableNameMap) {
        String tableName = tableNameMap.get(table.getName());
        if (StringUtils.isNotBlank(tableName)) {
            table.setName(tableName);
        }
        return table;
    }

    private static DynamicTableInfo extractInfo(Object paramObj) {
        // 判断参数里是否有DynamicTableInfo对象
        DynamicTableInfo info = null;
        if (paramObj instanceof DynamicTableInfo) {
            info = (DynamicTableInfo) paramObj;
        } else if (paramObj instanceof Map) {
            for (Object arg : ((Map<?, ?>) paramObj).values()) {
                if (arg instanceof DynamicTableInfo) {
                    info = (DynamicTableInfo) arg;
                    break;
                }
            }
        }
        return info;
    }

    private static String alterTableName(MappedStatement mappedStatement, Statement statement, Map<String, String> tableNameMap) {
        List<Table> tableList = new ArrayList<>();
        String originalSql = null;

        if (mappedStatement.getSqlCommandType() == SqlCommandType.INSERT) {
            Insert insertStatement = (Insert) statement;
            Table table = alterTableName(insertStatement.getTable(), tableNameMap);
            insertStatement.setTable(table);
            originalSql = insertStatement.toString();
        } else if (mappedStatement.getSqlCommandType() == SqlCommandType.DELETE) {
            Delete deleteStatement = (Delete) statement;
            Table table = alterTableName(deleteStatement.getTable(), tableNameMap);
            deleteStatement.setTable(table);
            originalSql = deleteStatement.toString();
        } else if (mappedStatement.getSqlCommandType() == SqlCommandType.UPDATE) {
            Update updateStatement = (Update) statement;
            for (Table p : updateStatement.getTables()) {
                Table table = alterTableName(p, tableNameMap);
                tableList.add(table);
            }
            updateStatement.setTables(tableList);
            originalSql = updateStatement.toString();
        } else if (mappedStatement.getSqlCommandType() == SqlCommandType.SELECT) {
            Select selectStatement = (Select) statement;
            PlainSelect plainSelect = (PlainSelect) selectStatement.getSelectBody();
            Table fromItem = (Table) plainSelect.getFromItem();
            Table table = alterTableName(fromItem, tableNameMap);
            plainSelect.setFromItem(table);
            List<Join> joinList = new ArrayList<>();
            if (!ObjectUtils.isEmpty(plainSelect.getJoins())) {
                for (Join p : plainSelect.getJoins()) {
                    if (p.getRightItem() != null) {
                        Table tableInner = alterTableName((Table) p.getRightItem(), tableNameMap);
                        p.setRightItem(tableInner);
                        joinList.add(p);
                    }
                }
                plainSelect.setJoins(joinList);
            }

            originalSql = plainSelect.toString();
        }

        return originalSql;
    }
}

@Configuration
public class MybatisPlusConfig {

    @Bean
    public DynamicTableInterceptor dynamicTableInterceptor(){
        return new DynamicTableInterceptor();
    }
}

public class TablePluginContainer {

    /**
     * 需要更改表名的缓存，请调用后手动清除
     * 当前表，需要替换的表名
     */
    public static ThreadLocal<Map<String,String>> TABLE_NAME_THREAD_LOCAL = new ThreadLocal<>();
}

@Data
public class DynamicTableInfo {

    private List<DynamicTableName> tableNames;
}


@Data
@AllArgsConstructor
public class DynamicTableName {

    /**
     * 当前表
     */
    private String currentTable;

    /**
     * 需要替换的表
     */
    private String itemTable;
}

void selectTest(DynamicTableInfo dynamicTableInfo);