feat: trace resolver mode (#89)

dbresolver.go

@@ -22,10 +22,11 @@ type DBResolver struct {
 }
 
 type Config struct {
-	Sources  []gorm.Dialector
-	Replicas []gorm.Dialector
-	Policy   Policy
-	datas    []interface{}
+	Sources           []gorm.Dialector
+	Replicas          []gorm.Dialector
+	Policy            Policy
+	datas             []interface{}
+	TraceResolverMode bool
 }
 
 func Register(config Config, datas ...interface{}) *DBResolver {
@@ -76,8 +77,9 @@ func (dr *DBResolver) compileConfig(config Config) (err error) {
 	var (
 		connPool = dr.DB.Config.ConnPool
 		r        = resolver{
-			dbResolver: dr,
-			policy:     config.Policy,
+			dbResolver:        dr,
+			policy:            config.Policy,
+			traceResolverMode: config.TraceResolverMode,
 		}
 	)
 
@@ -122,6 +124,10 @@ func (dr *DBResolver) compileConfig(config Config) (err error) {
 		}
 	}
 
+	if config.TraceResolverMode {
+		dr.Logger = NewResolverModeLogger(dr.Logger)
+	}
+
 	return nil
 }
 

dbresolver_test.go

@@ -2,10 +2,12 @@ package dbresolver_test
 
 import (
 	"fmt"
+	"os"
 	"testing"
 
 	"gorm.io/driver/mysql"
 	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
 	"gorm.io/plugin/dbresolver"
 )
 
@@ -50,16 +52,23 @@ func TestDBResolver(t *testing.T) {
 		if err != nil {
 			t.Fatalf("failed to connect db, got error: %v", err)
 		}
+		if debug := os.Getenv("DEBUG"); debug == "true" {
+			DB.Logger = DB.Logger.LogMode(logger.Info)
+		} else if debug == "false" {
+			DB.Logger = DB.Logger.LogMode(logger.Silent)
+		}
 
 		if err := DB.Use(dbresolver.Register(dbresolver.Config{
 			Sources: []gorm.Dialector{mysql.Open("gorm:gorm@tcp(localhost:9911)/gorm?charset=utf8&parseTime=True&loc=Local")},
 			Replicas: []gorm.Dialector{
 				mysql.Open("gorm:gorm@tcp(localhost:9912)/gorm?charset=utf8&parseTime=True&loc=Local"),
 				mysql.Open("gorm:gorm@tcp(localhost:9913)/gorm?charset=utf8&parseTime=True&loc=Local"),
 			},
+			TraceResolverMode: true,
 		}).Register(dbresolver.Config{
-			Sources:  []gorm.Dialector{mysql.Open("gorm:gorm@tcp(localhost:9914)/gorm?charset=utf8&parseTime=True&loc=Local")},
-			Replicas: []gorm.Dialector{mysql.Open("gorm:gorm@tcp(localhost:9913)/gorm?charset=utf8&parseTime=True&loc=Local")},
+			Sources:           []gorm.Dialector{mysql.Open("gorm:gorm@tcp(localhost:9914)/gorm?charset=utf8&parseTime=True&loc=Local")},
+			Replicas:          []gorm.Dialector{mysql.Open("gorm:gorm@tcp(localhost:9913)/gorm?charset=utf8&parseTime=True&loc=Local")},
+			TraceResolverMode: true,
 		}, "users", &Product{}).SetMaxOpenConns(5)); err != nil {
 			t.Fatalf("failed to use plugin, got error: %v", err)
 		}

logger.go

@@ -0,0 +1,54 @@
+package dbresolver
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+)
+
+type ResolverModeKey string
+type ResolverMode string
+
+const resolverModeKey ResolverModeKey = "dbresolver:resolver_mode_key"
+const (
+	ResolverModeSource  ResolverMode = "source"
+	ResolverModeReplica ResolverMode = "replica"
+)
+
+type resolverModeLogger struct {
+	logger.Interface
+}
+
+func (l resolverModeLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
+	var splitFn = func() (sql string, rowsAffected int64) {
+		sql, rowsAffected = fc()
+		op := ctx.Value(resolverModeKey)
+		if op != nil {
+			sql = fmt.Sprintf("[%s] %s", op, sql)
+			return
+		}
+
+		// the situation that dbresolver does not handle
+		// such as transactions, or some resolvers do not enable MarkResolverMode.
+		return
+	}
+	l.Interface.Trace(ctx, begin, splitFn, err)
+}
+
+func NewResolverModeLogger(l logger.Interface) logger.Interface {
+	if _, ok := l.(resolverModeLogger); ok {
+		return l
+	}
+	return resolverModeLogger{
+		Interface: l,
+	}
+}
+
+func markStmtResolverMode(stmt *gorm.Statement, mode ResolverMode) {
+	if _, ok := stmt.Logger.(resolverModeLogger); ok {
+		stmt.Context = context.WithValue(stmt.Context, resolverModeKey, mode)
+	}
+}

resolver.go

@@ -5,10 +5,11 @@ import (
 )
 
 type resolver struct {
-	sources    []gorm.ConnPool
-	replicas   []gorm.ConnPool
-	policy     Policy
-	dbResolver *DBResolver
+	sources           []gorm.ConnPool
+	replicas          []gorm.ConnPool
+	policy            Policy
+	dbResolver        *DBResolver
+	traceResolverMode bool
 }
 
 func (r *resolver) resolve(stmt *gorm.Statement, op Operation) (connPool gorm.ConnPool) {
@@ -18,10 +19,19 @@ func (r *resolver) resolve(stmt *gorm.Statement, op Operation) (connPool gorm.Co
 		} else {
 			connPool = r.policy.Resolve(r.replicas)
 		}
+		if r.traceResolverMode {
+			markStmtResolverMode(stmt, ResolverModeReplica)
+		}
 	} else if len(r.sources) == 1 {
 		connPool = r.sources[0]
+		if r.traceResolverMode {
+			markStmtResolverMode(stmt, ResolverModeSource)
+		}
 	} else {
 		connPool = r.policy.Resolve(r.sources)
+		if r.traceResolverMode {
+			markStmtResolverMode(stmt, ResolverModeSource)
+		}
 	}
 
 	if stmt.DB.PrepareStmt {