Log slow queries locally (#87)

* Log slow queries locally

Currently, queries slower than 100ms will be logged. This should not be
used in production.

* Fix settings functions and change default to 10ms

Author: thokra-nav

internal/cmd/api/api.go

@@ -78,8 +78,10 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error {
 	ctx, signalStop := signal.NotifyContext(ctx, syscall.SIGTERM, syscall.SIGINT)
 	defer signalStop()
 
+	dbSettings := []database.OptFunc{}
 	if cfg.WithFakeClients {
 		log.Warn("using fake clients")
+		dbSettings = append(dbSettings, database.WithSlowQueryLogger(10*time.Millisecond))
 	}
 
 	_, promReg, err := newMeterProvider(ctx)
@@ -88,7 +90,7 @@ func run(ctx context.Context, cfg *Config, log logrus.FieldLogger) error {
 	}
 
 	log.Info("connecting to database")
-	pool, err := database.New(ctx, cfg.DatabaseConnectionString, log.WithField("subsystem", "database"))
+	pool, err := database.New(ctx, cfg.DatabaseConnectionString, log.WithField("subsystem", "database"), dbSettings...)
 	if err != nil {
 		return fmt.Errorf("setting up database: %w", err)
 	}

internal/database/database.go

@@ -23,15 +23,34 @@ const databaseConnectRetries = 5
 
 var regParseSQLName = regexp.MustCompile(`\-\-\s*name:\s+(\S+)`)
 
-func New(ctx context.Context, dsn string, log logrus.FieldLogger) (*pgxpool.Pool, error) {
-	conn, err := NewPool(ctx, dsn, log, true)
+type settings struct {
+	slowQueryLogger time.Duration
+}
+
+type OptFunc func(*settings)
+
+// WithSlowQueryLogger enables slow query logging
+// This exposes attributes of the slow query to the logger, so it should not be used in production
+func WithSlowQueryLogger(d time.Duration) OptFunc {
+	return func(s *settings) {
+		s.slowQueryLogger = d
+	}
+}
+
+func New(ctx context.Context, dsn string, log logrus.FieldLogger, opts ...OptFunc) (*pgxpool.Pool, error) {
+	conn, err := NewPool(ctx, dsn, log, true, opts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect to the database: %w", err)
 	}
 	return conn, nil
 }
 
-func NewPool(ctx context.Context, dsn string, log logrus.FieldLogger, migrate bool) (*pgxpool.Pool, error) {
+func NewPool(ctx context.Context, dsn string, log logrus.FieldLogger, migrate bool, opts ...OptFunc) (*pgxpool.Pool, error) {
+	settings := &settings{}
+	for _, o := range opts {
+		o(settings)
+	}
+
 	if migrate {
 		if err := migrateDatabaseSchema("pgx", dsn, log); err != nil {
 			return nil, err
@@ -55,6 +74,14 @@ func NewPool(ctx context.Context, dsn string, log logrus.FieldLogger, migrate bo
 		}),
 	)
 
+	if settings.slowQueryLogger > 0 {
+		config.ConnConfig.Tracer = &slowQueryLogger{
+			log:      log,
+			sub:      config.ConnConfig.Tracer,
+			duration: settings.slowQueryLogger,
+		}
+	}
+
 	config.AfterConnect = func(ctx context.Context, conn *pgx.Conn) error {
 		t, err := conn.LoadType(ctx, "slug") // slug type
 		if err != nil {
@@ -113,3 +140,45 @@ func migrateDatabaseSchema(driver, dsn string, log logrus.FieldLogger) error {
 
 	return goose.Up(db, "migrations")
 }
+
+type slowQueryLogger struct {
+	log      logrus.FieldLogger
+	sub      pgx.QueryTracer
+	duration time.Duration
+}
+
+type sqlCtx int
+
+const sqlCtxKey sqlCtx = 0
+
+type bucket struct {
+	pgx.TraceQueryStartData
+	start time.Time
+}
+
+func (s *slowQueryLogger) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	ctx = s.sub.TraceQueryStart(ctx, conn, data)
+	return context.WithValue(ctx, sqlCtxKey, &bucket{
+		TraceQueryStartData: data,
+		start:               time.Now(),
+	})
+}
+
+func (s *slowQueryLogger) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+	s.sub.TraceQueryEnd(ctx, conn, data)
+
+	b, ok := ctx.Value(sqlCtxKey).(*bucket)
+	if !ok {
+		return
+	}
+
+	dur := time.Since(b.start)
+	if dur > s.duration {
+		s.log.WithFields(logrus.Fields{
+			"sql":  b.SQL,
+			"args": b.Args,
+			"time": dur,
+			"err":  data.Err,
+		}).Warn("slow query")
+	}
+}