MyBatis加载表统计

1.引入

统计当前应用, MyBatis所有使用到的表.

2.原理

1. 从Spring容器中, 依次获取 SqlSessionFactory, Configuration, MappedStatement
2. 获取SQL (mappedStatement.getBoundSql(params).getSql())
3. 解析SQL (Druid SQL-Parser)

3.源码

package com.clx.variablemanage;

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.conditions.query.QueryWrapper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.collections.MapUtils;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.util.Assert;

import java.util.*;
import java.util.stream.Collectors;

/**
 * MyBatis SQL 解析器
 *
 * @author liuxianqiang
 * @since 2021/11/19
 */
@Slf4j
public class MyBatisSqlAnalyzer {

    private static Map<TableStat.Name, TableStat> tableMap = new HashMap<>();

    public static void process(ApplicationContext context) {
        // 参数(mock)
        Map<String, Object> params = getMockParams();

        // 从Spring容器中获取 MappedStatement
        SqlSessionFactory sqlSessionFactory = context.getBean(SqlSessionFactory.class);
        Assert.notNull(sqlSessionFactory, "mybatis: sqlSessionFactory is null");

        Configuration configuration = sqlSessionFactory.getConfiguration();
        Assert.notNull(configuration, "mybatis: configuration is null");

        Collection<MappedStatement> mappedStatements = configuration.getMappedStatements();
        Assert.notEmpty(mappedStatements, "mybatis: mappedStatements is empty");

        // Statement去重 (sql会用namespace.funcId, funcId作为id, 注册两次)
        mappedStatements = mappedStatements.stream().distinct().collect(Collectors.toList());

        // 遍历MappedStatement, 解析SQL
        for (Object mappedStatement : mappedStatements) {
            if (!(mappedStatement instanceof MappedStatement)) {
                // ignore, 方法名ID的SQL重复
                continue;
            }
            parseStmt((MappedStatement)mappedStatement, params);
        }

        // 打印解析的信息
        log.info("---------------- table info ----------------");
        tableMap.forEach((k, v) -> {
            log.info("{}: {}, {}, {}, {}, {}",
                    k.getName(), v.getInsertCount(), v.getDeleteCount(), v.getUpdateCount(), v.getSelectCount(), v.getDropCount());
        });
    }

    private static Map<String, Object> getMockParams() {
        // 重写get方法
        Map<String, Object> params = new HashMap() {
            @Override
            public Object get(Object key) {
                Object v = super.get(key);
                return Objects.isNull(v) ? "1" : v;
            }
        };

        /**
         * 兼容mybatis-plus
         * @see  com.baomidou.mybatisplus.core.mapper.BaseMapper
         */
        params.put("coll", ImmutableList.of(1));
        params.put("et", new HashMap() {
            @Override
            public Object get(Object key) {
                return "1";
            }
        });
        params.put("cm", ImmutableMap.of("id", 1));
        params.put("ew", new QueryWrapper().eq("id", 1));
        params.put("MP_OPTLOCK_VERSION_COLUMN", "v");

        /**
         * 兼容项目
         */
        params.put("sessionIdList", ImmutableList.of(1));

        return params;
    }

    private static void parseStmt(MappedStatement mappedStatement, Map<String, Object> params) {
        String id = null;
        String sql = null;
        try {
            id = mappedStatement.getId();
            sql = mappedStatement.getBoundSql(params).getSql();

            List<SQLStatement> stmts = SQLUtils.parseStatements(sql, JdbcConstants.MYSQL);
            if (CollectionUtils.isNotEmpty(stmts)) {
                MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
                stmts.forEach(stmt -> {
                    stmt.accept(visitor);
                    appendTableState(visitor.getTables());
                });
            }
        } catch (Exception e) {
            log.error("parse sql exception, id:{}, sql:{}", id, sql, e);
        }
    }

    private static void appendTableState(Map<TableStat.Name, TableStat> tables) {
        if (MapUtils.isNotEmpty(tables)) {
            tables.forEach((k, v) -> {
                TableStat stat = tableMap.get(k);
                if (Objects.nonNull(stat)) {
                    v.setInsertCount(v.getInsertCount() + stat.getInsertCount());
                    v.setDeleteCount(v.getDeleteCount() + stat.getDeleteCount());
                    v.setUpdateCount(v.getUpdateCount() + stat.getUpdateCount());
                    v.setSelectCount(v.getSelectCount() + stat.getSelectCount());
                    v.setDropCount(v.getDropCount() + stat.getDropCount());
                }
                tableMap.put(k, v);
            });
        }
    }
}

4. 参考

• 相关sql解析器, antlr, apache calcite, druid sql parser, sqlparser
• Druid SQL-Parser
• 利用 druid 的 sql parser 模块解析 sql 语句获得表名