Skip to content

Commit

Permalink
Merge pull request #2098 from stringintech/tracelog-time-key
Browse files Browse the repository at this point in the history 
add TraceLogConfig with customizable TimeKey and ensureConfig method for default initialization
  • Loading branch information
@jackc
jackc authored Aug 13, 2024
2 parents c457de6 + 216049c commit d6fc8b0
Showing 1 changed file with 42 additions and 17 deletions.
59 changes: 42 additions & 17 deletions tracelog/tracelog.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()})
}
}

Expand Down Expand Up @@ -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})
}
}

Expand All @@ -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()})
}
}

Expand All @@ -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()
Expand All @@ -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
Expand All @@ -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,
})
}
}
Expand All @@ -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})
}
}

Expand Down

0 comments on commit d6fc8b0

Please sign in to comment.