原始需求
在SQL语句前面加上一个request-id
问题描述
今天收到业务同学反馈,说接入某个SDK后,request-id本地debug发现sql已经修改了,但打印的sql中却没有request-id信息
看了下代码,发现用户的代码其实就是下方 方案一代码,取不到的原因也在代码中注释了
方案一
package com.jiankunking.mybatisplus;
import static org.apache.commons.lang3.StringUtils.isNotBlank;
import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
import com.jiankunking.utils.JkkLogUtil;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.Properties;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
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.mapping.SqlCommandType;
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;
@Intercepts(
{
@Signature(type = StatementHandler.class, method = "prepare",
args = {Connection.class, Integer.class}),
@Signature(type = StatementHandler.class, method = "getBoundSql", args = {}),
@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}),
}
)
@Slf4j
public class MybatisPlusRequestIdInterceptor extends MybatisPlusInterceptor {
private Boolean enabled = true;
@Override
public Object intercept(Invocation invocation) throws Throwable {
Object target = invocation.getTarget();
Object[] args = invocation.getArgs();
BoundSql boundSql = null;
if (target instanceof Executor) {
Object parameter = args[1];
boolean isUpdate = args.length == 2;
MappedStatement ms = (MappedStatement) args[0];
if (!isUpdate && ms.getSqlCommandType() == SqlCommandType.SELECT) {
RowBounds rowBounds = (RowBounds) args[2];
if (args.length == 4) {
boundSql = ms.getBoundSql(parameter);
} else {
// 几乎不可能走进这里面,除非使用Executor的代理对象调用query[args[6]]
boundSql = (BoundSql) args[5];
}
}
} else {
StatementHandler statementHandler = (StatementHandler) target;
boundSql = statementHandler.getBoundSql();
}
if (enabled && boundSql != null) {
String sql = boundSql.getSql();
String requestId = JkkLogUtil.getCurrentRequestId();
StringBuilder sb = new StringBuilder("/*");
if (isNotBlank(requestId) && isNotBlank(sql)) {
sb.append(" Jkk-request-id:").append(requestId).append(" ");
}
String armsTraceId = JkkLogUtil.getTraceId();
if (isNotBlank(armsTraceId) && isNotBlank(sql)) {
sb.append(" trace-id:").append(armsTraceId).append(" ");
}
String armsRpcId = JkkLogUtil.getRpcId();
if (isNotBlank(armsRpcId)) {
sb.append(" rpc-id:").append(armsRpcId).append(" ");
}
sb.append(" */");
sb.append(sql);
// 通过反射修改sql语句
Field field = boundSql.getClass().getDeclaredField("sql");
field.setAccessible(true);
field.set(boundSql, sb.toString());
// 这里如果不调用 return invocation.proceed();
// 而是执行父类intercept的话,会导致反射修改的SQL无效
// 因为父类获取sql还是基于参数再次拼接的没有直接使用boundSql
}
return super.intercept(invocation);
}
@Override
public Object plugin(Object target) {
return Plugin.wrap(target, this);
}
@Override
public void setProperties(Properties properties) {
enabled = Boolean.valueOf(properties.getProperty("enabled", "true"));
}
}
方案二
直接修改mybatisplus的Statement
package com.jiankunking.mybatisplus;
import static org.apache.commons.lang3.StringUtils.isNotBlank;
import com.baomidou.mybatisplus.extension.plugins.MybatisPlusInterceptor;
import com.jiankunking.utils.JkkLogUtil;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.*;
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.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.reflection.factory.DefaultObjectFactory;
import org.apache.ibatis.reflection.wrapper.DefaultObjectWrapperFactory;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
@Intercepts(
{
@Signature(type = StatementHandler.class, method = "prepare",
args = {Connection.class, Integer.class}),
@Signature(type = StatementHandler.class, method = "getBoundSql", args = {}),
@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}),
}
)
@Slf4j
public class MybatisPlusRequestIdInterceptor extends MybatisPlusInterceptor {
private Boolean enabled = true;
@Override
public Object intercept(Invocation invocation) throws Throwable {
Object target = invocation.getTarget();
Object[] args = invocation.getArgs();
BoundSql boundSql = null;
if (target instanceof Executor) {
Object parameter = args[1];
boolean isUpdate = args.length == 2;
MappedStatement ms = (MappedStatement) args[0];
if (!isUpdate && ms.getSqlCommandType() == SqlCommandType.SELECT) {
RowBounds rowBounds = (RowBounds) args[2];
if (args.length == 4) {
boundSql = ms.getBoundSql(parameter);
} else {
// 几乎不可能走进这里面,除非使用Executor的代理对象调用query[args[6]]
boundSql = (BoundSql) args[5];
}
}
} else {
StatementHandler statementHandler = (StatementHandler) target;
boundSql = statementHandler.getBoundSql();
}
if (enabled && boundSql != null) {
String sql = boundSql.getSql();
String requestId = JkkLogUtil.getCurrentRequestId();
StringBuilder sb = new StringBuilder("/*");
if (isNotBlank(requestId) && isNotBlank(sql)) {
sb.append(" Jkk-request-id:").append(requestId).append(" ");
}
String armsTraceId = JkkLogUtil.getTraceId();
if (isNotBlank(armsTraceId) && isNotBlank(sql)) {
sb.append(" trace-id:").append(armsTraceId).append(" ");
}
String armsRpcId = JkkLogUtil.getRpcId();
if (isNotBlank(armsRpcId)) {
sb.append(" rpc-id:").append(armsRpcId).append(" ");
}
sb.append(" */");
sb.append(sql);
// 通过反射修改sql语句
//Field field = boundSql.getClass().getDeclaredField("sql");
//field.setAccessible(true);
//field.set(boundSql, sb.toString());
if (target instanceof Executor) {
setCurrentSql(invocation, sb.toString());
} else {
StatementHandler statementHandler = (StatementHandler) target;
MetaObject metaObject = SystemMetaObject.forObject(statementHandler);
metaObject.setValue("delegate.boundSql.sql", sb.toString());
}
}
return super.intercept(invocation);
}
@Override
public Object plugin(Object target) {
return Plugin.wrap(target, this);
}
@Override
public void setProperties(Properties properties) {
enabled = Boolean.valueOf(properties.getProperty("enabled", "true"));
}
private static final int MAPPED_STATEMENT_INDEX = 0;
private static final int PARAM_OBJ_INDEX = 1;
private void setCurrentSql(Invocation invocation, String sql) {
Object[] args = invocation.getArgs();
Object paramObj = args[PARAM_OBJ_INDEX];
MappedStatement mappedStatement = (MappedStatement) args[MAPPED_STATEMENT_INDEX];
BoundSql boundSql = mappedStatement.getBoundSql(paramObj);
MappedStatement newMappedStatement = copyFromMappedStatement(
mappedStatement, new BoundSqlSqlSource(boundSql));
MetaObject metaObject = MetaObject.forObject(newMappedStatement,
new DefaultObjectFactory(), new DefaultObjectWrapperFactory(),
new DefaultReflectorFactory());
metaObject.setValue("sqlSource.boundSql.sql", sql);
args[MAPPED_STATEMENT_INDEX] = newMappedStatement;
}
private class BoundSqlSqlSource implements SqlSource {
BoundSql boundSql;
public BoundSqlSqlSource(BoundSql boundSql) {
this.boundSql = boundSql;
}
public BoundSql getBoundSql(Object parameterObject) {
return boundSql;
}
}
@SuppressWarnings({"unchecked", "rawtypes"})
private MappedStatement copyFromMappedStatement(MappedStatement ms,
SqlSource newSqlSource) {
MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms
.getId(), newSqlSource, ms.getSqlCommandType());
builder.resource(ms.getResource());
builder.fetchSize(ms.getFetchSize());
builder.statementType(ms.getStatementType());
builder.keyGenerator(ms.getKeyGenerator());
// setStatementTimeout()
builder.timeout(ms.getTimeout());
// setParameterMap()
builder.parameterMap(ms.getParameterMap());
// setStatementResultMap()
List<ResultMap> resultMaps = new ArrayList<ResultMap>();
String id = "-inline";
if (ms.getResultMaps() != null) {
id = ms.getResultMaps().get(0).getId() + "-inline";
}
ResultMap resultMap = new ResultMap.Builder(null, id, Long.class,
new ArrayList()).build();
resultMaps.add(resultMap);
builder.resultMaps(resultMaps);
builder.resultSetType(ms.getResultSetType());
// setStatementCache()
builder.cache(ms.getCache());
builder.flushCacheRequired(ms.isFlushCacheRequired());
builder.useCache(ms.isUseCache());
return builder.build();
}
//private MappedStatement getMappedStatement(Invocation invo) {
// Object[] args = invo.getArgs();
// Object mappedStatement = args[MAPPED_STATEMENT_INDEX];
// return (MappedStatement) mappedStatement;
//}
}