Merge pull request #2098 from stringintech/tracelog-time-key

add TraceLogConfig with customizable TimeKey and ensureConfig method for default initialization

Author: jackc

tracelog/tracelog.go

@@ -117,11 +117,31 @@ func logQueryArgs(args []any) []any {
 	return logArgs
 }
 
+// TraceLogConfig holds the configuration for key names
+type TraceLogConfig struct {
+	TimeKey string
+}
+
+// DefaultTraceLogConfig returns the default configuration for TraceLog
+func DefaultTraceLogConfig() *TraceLogConfig {
+	return &TraceLogConfig{
+		TimeKey: "time",
+	}
+}
+
 // TraceLog implements pgx.QueryTracer, pgx.BatchTracer, pgx.ConnectTracer, and pgx.CopyFromTracer. All fields are
 // required.
 type TraceLog struct {
 	Logger   Logger
 	LogLevel LogLevel
+	Config   *TraceLogConfig
+}
+
+// ensureConfig initializes the Config field with default values if it is nil.
+func (tl *TraceLog) ensureConfig() {
+	if tl.Config == nil {
+		tl.Config = DefaultTraceLogConfig()
+	}
 }
 
 type ctxKey int
@@ -150,20 +170,21 @@ func (tl *TraceLog) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pg
 }
 
 func (tl *TraceLog) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+	tl.ensureConfig()
 	queryData := ctx.Value(tracelogQueryCtxKey).(*traceQueryData)
 
 	endTime := time.Now()
 	interval := endTime.Sub(queryData.startTime)
 
 	if data.Err != nil {
 		if tl.shouldLog(LogLevelError) {
-			tl.log(ctx, conn, LogLevelError, "Query", map[string]any{"sql": queryData.sql, "args": logQueryArgs(queryData.args), "err": data.Err, "time": interval})
+			tl.log(ctx, conn, LogLevelError, "Query", map[string]any{"sql": queryData.sql, "args": logQueryArgs(queryData.args), "err": data.Err, tl.Config.TimeKey: interval})
 		}
 		return
 	}
 
 	if tl.shouldLog(LogLevelInfo) {
-		tl.log(ctx, conn, LogLevelInfo, "Query", map[string]any{"sql": queryData.sql, "args": logQueryArgs(queryData.args), "time": interval, "commandTag": data.CommandTag.String()})
+		tl.log(ctx, conn, LogLevelInfo, "Query", map[string]any{"sql": queryData.sql, "args": logQueryArgs(queryData.args), tl.Config.TimeKey: interval, "commandTag": data.CommandTag.String()})
 	}
 }
 
@@ -191,20 +212,21 @@ func (tl *TraceLog) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pg
 }
 
 func (tl *TraceLog) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
+	tl.ensureConfig()
 	queryData := ctx.Value(tracelogBatchCtxKey).(*traceBatchData)
 
 	endTime := time.Now()
 	interval := endTime.Sub(queryData.startTime)
 
 	if data.Err != nil {
 		if tl.shouldLog(LogLevelError) {
-			tl.log(ctx, conn, LogLevelError, "BatchClose", map[string]any{"err": data.Err, "time": interval})
+			tl.log(ctx, conn, LogLevelError, "BatchClose", map[string]any{"err": data.Err, tl.Config.TimeKey: interval})
 		}
 		return
 	}
 
 	if tl.shouldLog(LogLevelInfo) {
-		tl.log(ctx, conn, LogLevelInfo, "BatchClose", map[string]any{"time": interval})
+		tl.log(ctx, conn, LogLevelInfo, "BatchClose", map[string]any{tl.Config.TimeKey: interval})
 	}
 }
 
@@ -223,20 +245,21 @@ func (tl *TraceLog) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data
 }
 
 func (tl *TraceLog) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
+	tl.ensureConfig()
 	copyFromData := ctx.Value(tracelogCopyFromCtxKey).(*traceCopyFromData)
 
 	endTime := time.Now()
 	interval := endTime.Sub(copyFromData.startTime)
 
 	if data.Err != nil {
 		if tl.shouldLog(LogLevelError) {
-			tl.log(ctx, conn, LogLevelError, "CopyFrom", map[string]any{"tableName": copyFromData.TableName, "columnNames": copyFromData.ColumnNames, "err": data.Err, "time": interval})
+			tl.log(ctx, conn, LogLevelError, "CopyFrom", map[string]any{"tableName": copyFromData.TableName, "columnNames": copyFromData.ColumnNames, "err": data.Err, tl.Config.TimeKey: interval})
 		}
 		return
 	}
 
 	if tl.shouldLog(LogLevelInfo) {
-		tl.log(ctx, conn, LogLevelInfo, "CopyFrom", map[string]any{"tableName": copyFromData.TableName, "columnNames": copyFromData.ColumnNames, "err": data.Err, "time": interval, "rowCount": data.CommandTag.RowsAffected()})
+		tl.log(ctx, conn, LogLevelInfo, "CopyFrom", map[string]any{"tableName": copyFromData.TableName, "columnNames": copyFromData.ColumnNames, "err": data.Err, tl.Config.TimeKey: interval, "rowCount": data.CommandTag.RowsAffected()})
 	}
 }
 
@@ -253,6 +276,7 @@ func (tl *TraceLog) TraceConnectStart(ctx context.Context, data pgx.TraceConnect
 }
 
 func (tl *TraceLog) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
+	tl.ensureConfig()
 	connectData := ctx.Value(tracelogConnectCtxKey).(*traceConnectData)
 
 	endTime := time.Now()
@@ -261,11 +285,11 @@ func (tl *TraceLog) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEn
 	if data.Err != nil {
 		if tl.shouldLog(LogLevelError) {
 			tl.Logger.Log(ctx, LogLevelError, "Connect", map[string]any{
-				"host":     connectData.connConfig.Host,
-				"port":     connectData.connConfig.Port,
-				"database": connectData.connConfig.Database,
-				"time":     interval,
-				"err":      data.Err,
+				"host":            connectData.connConfig.Host,
+				"port":            connectData.connConfig.Port,
+				"database":        connectData.connConfig.Database,
+				tl.Config.TimeKey: interval,
+				"err":             data.Err,
 			})
 		}
 		return
@@ -274,10 +298,10 @@ func (tl *TraceLog) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEn
 	if data.Conn != nil {
 		if tl.shouldLog(LogLevelInfo) {
 			tl.log(ctx, data.Conn, LogLevelInfo, "Connect", map[string]any{
-				"host":     connectData.connConfig.Host,
-				"port":     connectData.connConfig.Port,
-				"database": connectData.connConfig.Database,
-				"time":     interval,
+				"host":            connectData.connConfig.Host,
+				"port":            connectData.connConfig.Port,
+				"database":        connectData.connConfig.Database,
+				tl.Config.TimeKey: interval,
 			})
 		}
 	}
@@ -298,20 +322,21 @@ func (tl *TraceLog) TracePrepareStart(ctx context.Context, _ *pgx.Conn, data pgx
 }
 
 func (tl *TraceLog) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
+	tl.ensureConfig()
 	prepareData := ctx.Value(tracelogPrepareCtxKey).(*tracePrepareData)
 
 	endTime := time.Now()
 	interval := endTime.Sub(prepareData.startTime)
 
 	if data.Err != nil {
 		if tl.shouldLog(LogLevelError) {
-			tl.log(ctx, conn, LogLevelError, "Prepare", map[string]any{"name": prepareData.name, "sql": prepareData.sql, "err": data.Err, "time": interval})
+			tl.log(ctx, conn, LogLevelError, "Prepare", map[string]any{"name": prepareData.name, "sql": prepareData.sql, "err": data.Err, tl.Config.TimeKey: interval})
 		}
 		return
 	}
 
 	if tl.shouldLog(LogLevelInfo) {
-		tl.log(ctx, conn, LogLevelInfo, "Prepare", map[string]any{"name": prepareData.name, "sql": prepareData.sql, "time": interval, "alreadyPrepared": data.AlreadyPrepared})
+		tl.log(ctx, conn, LogLevelInfo, "Prepare", map[string]any{"name": prepareData.name, "sql": prepareData.sql, tl.Config.TimeKey: interval, "alreadyPrepared": data.AlreadyPrepared})
 	}
 }