From 216049c62beb2a91e34caa197c7ea3223be33a6f Mon Sep 17 00:00:00 2001 From: stringintech Date: Fri, 26 Jul 2024 22:58:14 +0330 Subject: [PATCH] add TraceLogConfig with customizable TimeKey and ensureConfig method for default initialization --- tracelog/tracelog.go | 59 +++++++++++++++++++++++++++++++------------- 1 file changed, 42 insertions(+), 17 deletions(-) diff --git a/tracelog/tracelog.go b/tracelog/tracelog.go index 78b15fbe4..212a25407 100644 --- a/tracelog/tracelog.go +++ b/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,6 +170,7 @@ 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() @@ -157,13 +178,13 @@ func (tl *TraceLog) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx. 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,6 +212,7 @@ 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() @@ -198,13 +220,13 @@ func (tl *TraceLog) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx. 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,6 +245,7 @@ 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() @@ -230,13 +253,13 @@ func (tl *TraceLog) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data p 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,6 +322,7 @@ 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() @@ -305,13 +330,13 @@ func (tl *TraceLog) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pg 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}) } }