mybaits实现多租户动态数据源

mybaits实现多租户动态数据源

1.添加mybatis拦截器，拦截mybatis执行sql
2.配置不需要的动态切换的表名filterTableNames
3.使用druid解析sql，获取sql中的全部表名getTableNameBySql
4.定义获取动态数据库的方法getSlaveName，可通过request中的参数或用户信息，进行自定义切换对应的数据库
5.在拦截器中拦截sql，替换sql语句中的表名，增加数据库前缀，达到实现动态数据源

import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.dialect.mysql.visitor.MySqlSchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;
import com.alibaba.druid.util.JdbcConstants;
import com.baomidou.mybatisplus.core.toolkit.PluginUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import javax.servlet.http.HttpServletRequest;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;

/**
 * 多租户数据源动态切换拦截器
 *
 * @author Enzo
 * @date 2022-12-2
 **/
@Slf4j
@Component
@Intercepts({@Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})})
public class SwitchDbInterceptor implements Interceptor {

    private static final String MODIFY_SQL_PATH = "parameterHandler.boundSql.sql";
    private static final String HEADER_SWITCH_DB = "switchDb";
    /**
     * 忽略切换数据源的表名
     **/
    private static List<String> filterTableNames;

    static {
        filterTableNames = new ArrayList<>();
        filterTableNames.add("DYNAMIC_PROJECT_SPACE");
        filterTableNames.add("DYNAMIC_PROJECT_SPACE_USER");
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        //获取需要切换的数据库名
        String switchDBName = getSlaveName();
        if (StringUtils.isNotEmpty(switchDBName)) {
            StatementHandler statementHandler = PluginUtils.realTarget(invocation.getTarget());
            MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
            String originalSql = (String) metaObject.getValue(MODIFY_SQL_PATH);
            List<String> tableNameList = null;
            try {
                //获取sql中的全部表名
                tableNameList = getTableNameBySql(originalSql);
            } catch (Exception e) {
                log.error(e.toString(), e);
            }
            if (tableNameList != null) {
                //更新sql语句，表名加上数据库前缀
                String sql = originalSql.replace("\n", " ").replace("\t", " ");
                for (String tableName : tableNameList) {
                    //忽略的动态切换数据库的表名
                    if (!filterTableNames.contains(tableName.toUpperCase())) {
                        sql = sql.replace(" " + tableName + " ", " " + switchDBName + "." + tableName + " ");
                    }
                }
                metaObject.setValue("delegate.boundSql.sql", sql);
            }
        }
        return invocation.proceed();
    }

    /**
     * 使用druid获取sql中的表名
     *
     * @param sql
     * @return java.util.List<java.lang.String>
     * @author Enzo
     * @date 2022-12-2
     **/
    private static List<String> getTableNameBySql(String sql) {
        String dbType = JdbcConstants.MYSQL;
        List<String> tableNameList = new ArrayList<>();
        try {
            List<SQLStatement> stmtList = SQLUtils.parseStatements(sql, dbType);
            if (CollectionUtils.isNotEmpty(stmtList)) {
                for (SQLStatement sqlStatement : stmtList) {
                    MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
                    sqlStatement.accept(visitor);
                    Map<TableStat.Name, TableStat> tables = visitor.getTables();
                    Set<TableStat.Name> tableNameSet = tables.keySet();
                    for (TableStat.Name name : tableNameSet) {
                        String tableName = name.getName();
                        if (StringUtils.isNotBlank(tableName)) {
                            tableNameList.add(tableName);
                        }
                    }
                }
            }
        } catch (Exception e) {
            log.error("**************异常SQL:[{}]*****************\\n", sql);
            log.error(e.getMessage(), e);
        }
        return tableNameList;
    }


    /**
     * 根据请求参数或用户获取动态数据库名
     *
     * @return java.lang.String
     * @author Enzo
     * @date 2022-12-2
     **/
    private String getSlaveName() {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        String databaseName = null;
        if (attributes != null) {
            HttpServletRequest request = attributes.getRequest();
            String switchDb = request.getHeader(HEADER_SWITCH_DB);
            if (StringUtils.isNotEmpty(switchDb)) {
                databaseName = switchDb;
            }
        }
        return databaseName;
    }

}