From 503beee54124386012a3614ce4b0732444be21f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BC=80=E6=BA=90=E6=B5=B7=E5=93=A5?= <fuhai999@gmail.com>
Date: Wed, 12 Apr 2023 18:40:16 +0800
Subject: [PATCH] =?UTF-8?q?feature:=20SQL=E5=AE=A1=E8=AE=A1=E6=96=B0?=
 =?UTF-8?q?=E5=A2=9E=E6=95=B0=E6=8D=AE=E8=BF=94=E5=9B=9E=E8=A1=8C=E6=95=B0?=
 =?UTF-8?q?=20close=20#I6UZLJ?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../mybatisflex/core/audit/AuditManager.java  | 14 ++++++++----
 .../mybatisflex/core/audit/AuditMessage.java  | 22 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditManager.java b/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditManager.java
index 13a97b65..225d9f55 100644
--- a/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditManager.java
+++ b/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditManager.java
@@ -62,8 +62,8 @@ public class AuditManager {
     }
 
 
-    public static void setMessageReporter(MessageReporter messageReporter){
-        MessageCollector newMessageCollector = new ScheduledMessageCollector(10,messageReporter);
+    public static void setMessageReporter(MessageReporter messageReporter) {
+        MessageCollector newMessageCollector = new ScheduledMessageCollector(10, messageReporter);
         setMessageCollector(newMessageCollector);
     }
 
@@ -75,7 +75,7 @@ public class AuditManager {
     }
 
     private static void releaseScheduledMessageCollector(MessageCollector messageCollector) {
-        if (messageCollector instanceof ScheduledMessageCollector){
+        if (messageCollector instanceof ScheduledMessageCollector) {
             ((ScheduledMessageCollector) messageCollector).release();
         }
     }
@@ -87,7 +87,13 @@ public class AuditManager {
         }
         auditMessage.setQueryTime(clock.getTick());
         try {
-            return supplier.execute();
+            T result = supplier.execute();
+            if (result instanceof Collection) {
+                auditMessage.setQueryCount(((Collection) result).size());
+            } else if (result != null) {
+                auditMessage.setQueryCount(1);
+            }
+            return result;
         } finally {
             auditMessage.setElapsedTime(clock.getTick() - auditMessage.getQueryTime());
             auditMessage.setQuery(boundSql.getSql());
diff --git a/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditMessage.java b/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditMessage.java
index 6375c50f..b4f6ecdf 100644
--- a/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditMessage.java
+++ b/mybatis-flex-core/src/main/java/com/mybatisflex/core/audit/AuditMessage.java
@@ -32,6 +32,7 @@ public class AuditMessage implements Serializable {
     private String platform;
     private String module;
     private String url;
+    private String bizId; //自定义业务ID
 
     private String user;
     private String userIp;
@@ -39,6 +40,7 @@ public class AuditMessage implements Serializable {
 
     private String query;
     private List<Object> queryParams;
+    private int queryCount;
 
     private long queryTime;     // Sql 执行的当前时间，单位毫秒
     private long elapsedTime;   // Sql 执行消耗的时间，单位毫秒
@@ -70,6 +72,14 @@ public class AuditMessage implements Serializable {
         this.url = url;
     }
 
+    public String getBizId() {
+        return bizId;
+    }
+
+    public void setBizId(String bizId) {
+        this.bizId = bizId;
+    }
+
     public String getUser() {
         return user;
     }
@@ -148,6 +158,13 @@ public class AuditMessage implements Serializable {
                 });
     }
 
+    public int getQueryCount() {
+        return queryCount;
+    }
+
+    public void setQueryCount(int queryCount) {
+        this.queryCount = queryCount;
+    }
 
     public long getQueryTime() {
         return queryTime;
@@ -182,17 +199,20 @@ public class AuditMessage implements Serializable {
 
     @Override
     public String toString() {
-        return "{" +
+        return "AuditMessage{" +
                 "platform='" + platform + '\'' +
                 ", module='" + module + '\'' +
                 ", url='" + url + '\'' +
+                ", bizId='" + bizId + '\'' +
                 ", user='" + user + '\'' +
                 ", userIp='" + userIp + '\'' +
                 ", hostIp='" + hostIp + '\'' +
                 ", query='" + query + '\'' +
                 ", queryParams=" + queryParams +
+                ", queryCount=" + queryCount +
                 ", queryTime=" + queryTime +
                 ", elapsedTime=" + elapsedTime +
+                ", metas=" + metas +
                 '}';
     }
 }