diff --git a/callbacks.go b/callbacks.go index 83d103df..fdde21e9 100644 --- a/callbacks.go +++ b/callbacks.go @@ -74,7 +74,6 @@ func (cs *callbacks) Raw() *processor { func (p *processor) Execute(db *DB) { curTime := time.Now() stmt := db.Statement - db.RowsAffected = 0 if stmt.Model == nil { stmt.Model = stmt.Dest diff --git a/callbacks/row.go b/callbacks/row.go index a36c0116..4f985d7b 100644 --- a/callbacks/row.go +++ b/callbacks/row.go @@ -16,6 +16,8 @@ func RowQuery(db *gorm.DB) { } else { db.Statement.Dest = db.Statement.ConnPool.QueryRowContext(db.Statement.Context, db.Statement.SQL.String(), db.Statement.Vars...) } + + db.RowsAffected = -1 } } } diff --git a/finisher_api.go b/finisher_api.go index 2c56d763..63061553 100644 --- a/finisher_api.go +++ b/finisher_api.go @@ -8,6 +8,7 @@ import ( "strings" "gorm.io/gorm/clause" + "gorm.io/gorm/logger" "gorm.io/gorm/schema" "gorm.io/gorm/utils" ) @@ -353,7 +354,9 @@ func (db *DB) Rows() (*sql.Rows, error) { // Scan scan value to a struct func (db *DB) Scan(dest interface{}) (tx *DB) { + currentLogger, newLogger := db.Logger, logger.Recorder.New() tx = db.getInstance() + tx.Logger = newLogger if rows, err := tx.Rows(); err != nil { tx.AddError(err) } else { @@ -362,6 +365,11 @@ func (db *DB) Scan(dest interface{}) (tx *DB) { tx.ScanRows(rows, dest) } } + + currentLogger.Trace(tx.Statement.Context, newLogger.BeginAt, func() (string, int64) { + return newLogger.SQL, tx.RowsAffected + }, tx.Error) + tx.Logger = currentLogger return } diff --git a/logger/logger.go b/logger/logger.go index e568fb24..b278ad5d 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -63,6 +63,7 @@ var ( LogLevel: Warn, Colorful: true, }) + Recorder = traceRecorder{Interface: Default} ) func New(writer Writer, config Config) Interface { @@ -70,18 +71,18 @@ func New(writer Writer, config Config) Interface { infoStr = "%s\n[info] " warnStr = "%s\n[warn] " errStr = "%s\n[error] " - traceStr = "%s\n[%.3fms] [rows:%d] %s" - traceWarnStr = "%s\n[%.3fms] [rows:%d] %s" - traceErrStr = "%s %s\n[%.3fms] [rows:%d] %s" + traceStr = "%s\n[%.3fms] [rows:%v] %s" + traceWarnStr = "%s\n[%.3fms] [rows:%v] %s" + traceErrStr = "%s %s\n[%.3fms] [rows:%v] %s" ) if config.Colorful { infoStr = Green + "%s\n" + Reset + Green + "[info] " + Reset warnStr = BlueBold + "%s\n" + Reset + Magenta + "[warn] " + Reset errStr = Magenta + "%s\n" + Reset + Red + "[error] " + Reset - traceStr = Green + "%s\n" + Reset + Yellow + "[%.3fms] " + BlueBold + "[rows:%d]" + Reset + " %s" - traceWarnStr = Green + "%s\n" + Reset + RedBold + "[%.3fms] " + Yellow + "[rows:%d]" + Magenta + " %s" + Reset - traceErrStr = RedBold + "%s " + MagentaBold + "%s\n" + Reset + Yellow + "[%.3fms] " + BlueBold + "[rows:%d]" + Reset + " %s" + traceStr = Green + "%s\n" + Reset + Yellow + "[%.3fms] " + BlueBold + "[rows:%v]" + Reset + " %s" + traceWarnStr = Green + "%s\n" + Reset + RedBold + "[%.3fms] " + Yellow + "[rows:%v]" + Magenta + " %s" + Reset + traceErrStr = RedBold + "%s " + MagentaBold + "%s\n" + Reset + Yellow + "[%.3fms] " + BlueBold + "[rows:%v]" + Reset + " %s" } return &logger{ @@ -138,13 +139,43 @@ func (l logger) Trace(ctx context.Context, begin time.Time, fc func() (string, i switch { case err != nil && l.LogLevel >= Error: sql, rows := fc() - l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, rows, sql) + if rows == -1 { + l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, "-", sql) + } else { + l.Printf(l.traceErrStr, utils.FileWithLineNum(), err, float64(elapsed.Nanoseconds())/1e6, rows, sql) + } case elapsed > l.SlowThreshold && l.SlowThreshold != 0 && l.LogLevel >= Warn: sql, rows := fc() - l.Printf(l.traceWarnStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql) + if rows == -1 { + l.Printf(l.traceWarnStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, "-", sql) + } else { + l.Printf(l.traceWarnStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql) + } case l.LogLevel >= Info: sql, rows := fc() - l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql) + if rows == -1 { + l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, "-", sql) + } else { + l.Printf(l.traceStr, utils.FileWithLineNum(), float64(elapsed.Nanoseconds())/1e6, rows, sql) + } } } } + +type traceRecorder struct { + Interface + BeginAt time.Time + SQL string + RowsAffected int64 + Err error +} + +func (l traceRecorder) New() *traceRecorder { + return &traceRecorder{Interface: l.Interface} +} + +func (l *traceRecorder) Trace(ctx context.Context, begin time.Time, fc func() (string, int64), err error) { + l.BeginAt = begin + l.SQL, l.RowsAffected = fc() + l.Err = err +} diff --git a/scan.go b/scan.go index d7cddbe6..8d737b17 100644 --- a/scan.go +++ b/scan.go @@ -52,6 +52,7 @@ func scanIntoMap(mapValue map[string]interface{}, values []interface{}, columns func Scan(rows *sql.Rows, db *DB, initialized bool) { columns, _ := rows.Columns() values := make([]interface{}, len(columns)) + db.RowsAffected = 0 switch dest := db.Statement.Dest.(type) { case map[string]interface{}, *map[string]interface{}: