feat/sqlx_metric

zeromicro · Jan 19, 2025 · b9b19e3 · b9b19e3
1 parent b650c8c
commit b9b19e3
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/core/stores/sqlx/sqlmanager.go b/core/stores/sqlx/sqlmanager.go
@@ -1,10 +1,14 @@
 package sqlx
 
 import (
+	"crypto/sha256"
 	"database/sql"
+	"encoding/hex"
 	"io"
 	"time"
 
+	"github.com/go-sql-driver/mysql"
+
 	"github.com/zeromicro/go-zero/core/syncx"
 )
 
@@ -23,6 +27,21 @@ func getCachedSqlConn(driverName, server string) (*sql.DB, error) {
 			return nil, err
 		}
 
+		if driverName == mysqlDriverName {
+			if cfg, err := mysql.ParseDSN(server); err != nil {
+				// do nothing here
+			} else {
+				checksum := sha256.Sum256([]byte(server))
+				connCollector.registerClient(&statGetter{
+					dbName: cfg.DBName,
+					hash:   hex.EncodeToString(checksum[:]),
+					poolStats: func() sql.DBStats {
+						return conn.Stats()
+					},
+				})
+			}
+		}
+
 		return conn, nil
 	})
 	if err != nil {