Bluemage/go/pkg/log/sql.go

package log

import (
	"context"
	"log/slog"
	"strings"

	sqldblogger "github.com/simukti/sqldb-logger"
)

var _ sqldblogger.Logger = (*SQLLogger)(nil)

type SQLLogger struct{}

var _ slog.LogValuer = (*QueryCollector)(nil)

type QueryCollector struct {
	Statement string
	Args      any
}

func (qu *QueryCollector) LogValue() slog.Value {
	return slog.GroupValue(
		slog.String("statement", qu.Statement),
		slog.Any("args", qu.Args),
	)
}

type queryCollectorKey struct{}

func QueryCollectorFromContext(ctx context.Context) *QueryCollector {
	coll, _ := ctx.Value(queryCollectorKey{}).(*QueryCollector)
	return coll
}

func WithQueryCollector(ctx context.Context) (context.Context, *QueryCollector) {
	coll := &QueryCollector{}
	return context.WithValue(ctx, queryCollectorKey{}, coll), coll
}

func (SQLLogger) Log(ctx context.Context, level sqldblogger.Level, msg string, data map[string]interface{}) {
	if qc := QueryCollectorFromContext(ctx); qc != nil {
		qc.Statement = msg
		qc.Args = data["args"]
		return
	}
	slog.DebugContext(ctx, strings.ReplaceAll(msg, "\n", " "), "lvl", level.String(), "data", data)
}