mybatis日志plugin

package com.wujie.growth.awardcenter.repository.plugin;

import com.cxyx.common.db.config.DbConfig;
import com.cxyx.common.rpc.Context;
import com.cxyx.common.util.ContextUtil;
import com.cxyx.common.util.JsonUtil;
import com.cxyx.common.util.PropertyUtil;
import com.xiaoju.apollo.message.StringUtil;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.ibatis.binding.MapperMethod;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.apache.ibatis.session.defaults.DefaultSqlSession;
import org.springframework.stereotype.Component;
import java.io.StringReader;
import java.lang.reflect.Field;
import java.text.SimpleDateFormat;
import java.util.Collection;
import java.util.Date;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import static com.wujie.growth.awardcenter.common.constant.AwardCenterConstant.PRINT_MYSQL_LOG_THREAD_LOCAL;

/**

  • @author xuan

  • @create 2020-12-20 16:37
    **/
    @Setter
    @Slf4j
    @Intercepts({
    @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
    @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
    @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
    })
    @Component
    public class MybatisLogPlugin implements Interceptor {

    private static final CCJSqlParserManager PARSER_MANAGER = new CCJSqlParserManager();

    private DbConfig dbConfig = PropertyUtil.newInstance(DbConfig.class).orElse(new DbConfig());

    private static String FOREACH_PREFIX = "_frch";

    /**

    • 积攒过多就悄悄丢弃
      */
      private static ThreadPoolExecutor pool = new ThreadPoolExecutor(2, 4, 60, TimeUnit.SECONDS, new ArrayBlockingQueue<>(1000), new ThreadPoolExecutor.DiscardPolicy());

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
    long startTime = System.currentTimeMillis();
    boolean errorFlag = false;
    Object result = null;
    try {
    result = invocation.proceed();
    } catch (Throwable e) {
    errorFlag = true;
    result = e;
    throw e;
    } finally {
    try {
    //优先读取打印日志的后门标记(不区分增删改查-都打印):不要清理printMysqlLogFlag,否则一次request请求只会打印第一条SQL
    Boolean printMysqlLogFlag = PRINT_MYSQL_LOG_THREAD_LOCAL.get();
    if (printMysqlLogFlag != null && printMysqlLogFlag) {
    printSql(invocation, startTime, errorFlag, result);
    } else if ("update".equals(invocation.getMethod().getName())) {
    printSql(invocation, startTime, errorFlag, result);
    }
    } catch (Exception e) {
    log.error("mybatis-log-plugin error", e);
    }
    }
    return result;
    }

private void printSql(Invocation invocation, long startTime, boolean errorFlag, Object result) {
    Map contextParams = ContextUtil.getAll();
    pool.execute(() -> {
        ContextUtil.putAll(contextParams);
        Object parameter = invocation.getArgs()[1];
        BoundSql boundSql = ((MappedStatement) invocation.getArgs()[0]).getBoundSql(parameter);
        String sql = boundSql.getSql();
        // 格式化Sql语句,去除换行符,替换参数
        sql = formatSql(sql, boundSql);
        //分表时表名替换
        sql = shardingTableReplaceName(sql);
        if (errorFlag) {
            log.error("sql=[{}]||result={}||proc_time={}", sql, JsonUtil.toString(result), System.currentTimeMillis() - startTime);
        } else {
            log.info("sql=[{}]||result={}||proc_time={}", sql, JsonUtil.toString(result), System.currentTimeMillis() - startTime);
        }
    });
}

private String shardingTableReplaceName(String sql) {
    Set tableNames = getTableNames(sql);
    if (tableNames.size() == 1) {
        final String tableName = tableNames.iterator().next();
        String newTableName = tableName;
        if (tableNeedSharding(tableName)) {
            newTableName = setSubfixSharding(tableName);
        }
        //压测数据暂时不做处理
        if (!org.apache.commons.lang3.StringUtils.equals(tableName, newTableName)) {
            sql = replace(sql, tableName, newTableName);
        }
    }
    return sql;
}

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

@Override
public void setProperties(Properties properties) {
}

@SuppressWarnings("unchecked")
private String formatSql(String sql, BoundSql boundSql) {
    if (StringUtil.isBlank(sql)) {
        return "";
    }
    sql = beautifySql(sql);
    List parameterMappingList = boundSql.getParameterMappings();
    Object parameterObject = boundSql.getParameterObject();
    if (parameterObject == null || CollectionUtils.isEmpty(parameterMappingList)) {
        return sql;
    }
    String sqlWithoutReplacePlaceholder = sql;
    try {
        Class parameterObjectClass = parameterObject.getClass();
        // 如果参数是StrictMap且Value类型为Collection,获取key="list"的属性,这里主要是为了处理循环时传入List这种参数的占位符替换
        // 例如select * from xxx where id in ...
        if (isStrictMap(parameterObjectClass)) {
            DefaultSqlSession.StrictMap> strictMap = (DefaultSqlSession.StrictMap>) parameterObject;
            if (isList(strictMap.get("list").getClass())) {
                sql = handleListParameter(sql, strictMap.get("list"));
            }
        } else if (isMap(parameterObjectClass)) {
            // 如果参数是Map则直接强转,通过map.get(key)方法获取真正的属性值
            // 这里主要是为了处理