feat: add stmtId for AuditMessage

This commit is contained in:
pbnoyz 2025-08-26 22:36:36 +08:00
parent c65a3608f8
commit 9afcb9514d
5 changed files with 77 additions and 5 deletions

View File

@ -97,11 +97,12 @@ public class AuditManager {
}
@SuppressWarnings("rawtypes")
public static <T> T startAudit(AuditRunnable<T> supplier, Statement statement, BoundSql boundSql, Configuration configuration) throws SQLException {
public static <T> T startAudit(AuditRunnable<T> supplier, String stmtId, Statement statement, BoundSql boundSql, Configuration configuration) throws SQLException {
AuditMessage auditMessage = messageFactory.create();
if (auditMessage == null) {
return supplier.execute();
}
auditMessage.setStmtId(stmtId);
String key = DataSourceKey.get();
if (StringUtil.noText(key)) {
key = FlexGlobalConfig.getDefaultConfig()

View File

@ -95,6 +95,11 @@ public class AuditMessage implements Serializable {
*/
private long elapsedTime;
/**
* MappedStatement ID
*/
private String stmtId;
/**
* 数据库名称
*/
@ -270,6 +275,14 @@ public class AuditMessage implements Serializable {
metas.put(key, value);
}
public String getStmtId() {
return stmtId;
}
public void setStmtId(String stmtId) {
this.stmtId = stmtId;
}
public String getDsName() {
return dsName;
}
@ -293,6 +306,7 @@ public class AuditMessage implements Serializable {
", queryCount=" + queryCount +
", queryTime=" + queryTime +
", elapsedTime=" + elapsedTime +
", stmtId=" + stmtId +
", dsName=" + dsName +
", metas=" + metas +
'}';

View File

@ -46,9 +46,11 @@ public class FlexStatementHandler implements StatementHandler {
private final BoundSql boundSql;
private final boolean auditEnable = AuditManager.isAuditEnable();
private final Configuration configuration;
private final String stmtId;
public FlexStatementHandler(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
configuration = ms.getConfiguration();
stmtId = ms.getId();
switch (ms.getStatementType()) {
case STATEMENT:
delegate = new SimpleStatementHandler(executor, ms, parameter, rowBounds, resultHandler, boundSql);
@ -83,7 +85,7 @@ public class FlexStatementHandler implements StatementHandler {
AuditManager.startAudit(() -> {
delegate.batch(statement);
return null;
}, statement, boundSql, configuration);
}, stmtId, statement, boundSql, configuration);
} else {
delegate.batch(statement);
}
@ -91,19 +93,19 @@ public class FlexStatementHandler implements StatementHandler {
@Override
public int update(Statement statement) throws SQLException {
return auditEnable ? AuditManager.startAudit(() -> delegate.update(statement), statement, boundSql, configuration)
return auditEnable ? AuditManager.startAudit(() -> delegate.update(statement), stmtId, statement, boundSql, configuration)
: delegate.update(statement);
}
@Override
public <E> List<E> query(Statement statement, ResultHandler resultHandler) throws SQLException {
return auditEnable ? AuditManager.startAudit(() -> delegate.query(statement, resultHandler), statement, boundSql, configuration)
return auditEnable ? AuditManager.startAudit(() -> delegate.query(statement, resultHandler), stmtId, statement, boundSql, configuration)
: delegate.query(statement, resultHandler);
}
@Override
public <E> Cursor<E> queryCursor(Statement statement) throws SQLException {
return auditEnable ? AuditManager.startAudit(() -> delegate.queryCursor(statement), statement, boundSql, configuration)
return auditEnable ? AuditManager.startAudit(() -> delegate.queryCursor(statement), stmtId, statement, boundSql, configuration)
: delegate.queryCursor(statement);
}

View File

@ -18,6 +18,9 @@ package com.mybatisflex.test.mapper;
import com.mybatisflex.core.BaseMapper;
import com.mybatisflex.test.model.Account;
import org.apache.ibatis.annotations.Select;
import java.util.List;
/**
* @author 王帅
@ -25,4 +28,7 @@ import com.mybatisflex.test.model.Account;
*/
public interface AccountMapper extends BaseMapper<Account> {
@Select("select * from tb_account")
List<Account> selectAccounts();
}

View File

@ -0,0 +1,49 @@
package com.mybatisflex.test;
import com.mybatisflex.core.audit.AuditManager;
import com.mybatisflex.core.audit.AuditMessage;
import com.mybatisflex.core.audit.MessageCollector;
import com.mybatisflex.test.mapper.AccountMapper;
import com.mybatisflex.test.mapper.TbClassMapper;
import lombok.Getter;
import org.assertj.core.api.WithAssertions;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
import java.util.ArrayList;
import java.util.List;
@RunWith(SpringJUnit4ClassRunner.class)
@ContextConfiguration(classes = AppConfig.class)
public class AuditTest implements WithAssertions {
@Autowired
AccountMapper accountMapper;
@Autowired
TbClassMapper tbClassMapper;
@Test
public void test() {
TestMsgCollector collector = new TestMsgCollector();
AuditManager.setMessageCollector(collector);
AuditManager.setAuditEnable(true);
accountMapper.selectAccounts();
tbClassMapper.selectAll();
List<AuditMessage> messages = collector.getMessages();
assertThat(messages.size()).isEqualTo(2);
assertThat(messages.get(0).getStmtId()).isEqualTo("com.mybatisflex.test.mapper.AccountMapper.selectAccounts");
assertThat(messages.get(1).getStmtId()).isEqualTo("com.mybatisflex.test.mapper.TbClassMapper.selectListByQuery");
}
static class TestMsgCollector implements MessageCollector {
@Getter
private final List<AuditMessage> messages = new ArrayList<>();
@Override
public void collect(AuditMessage message) {
messages.add(message);
}
}
}