automatically ignore known forwarded addresses, fixes #64

vendor/go.mau.fi/util/dbutil/connlog.go

@@ -0,0 +1,189 @@
+// Copyright (c) 2022 Tulir Asokan
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package dbutil
+
+import (
+	"context"
+	"database/sql"
+	"time"
+)
+
+// LoggingExecable is a wrapper for anything with database Exec methods (i.e. sql.Conn, sql.DB and sql.Tx)
+// that can preprocess queries (e.g. replacing $ with ? on SQLite) and log query durations.
+type LoggingExecable struct {
+	UnderlyingExecable UnderlyingExecable
+	db                 *Database
+}
+
+func (le *LoggingExecable) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	start := time.Now()
+	query = le.db.mutateQuery(query)
+	res, err := le.UnderlyingExecable.ExecContext(ctx, query, args...)
+	le.db.Log.QueryTiming(ctx, "Exec", query, args, -1, time.Since(start), err)
+	return res, err
+}
+
+func (le *LoggingExecable) QueryContext(ctx context.Context, query string, args ...interface{}) (Rows, error) {
+	start := time.Now()
+	query = le.db.mutateQuery(query)
+	rows, err := le.UnderlyingExecable.QueryContext(ctx, query, args...)
+	le.db.Log.QueryTiming(ctx, "Query", query, args, -1, time.Since(start), err)
+	return &LoggingRows{
+		ctx:   ctx,
+		db:    le.db,
+		query: query,
+		args:  args,
+		rs:    rows,
+		start: start,
+	}, err
+}
+
+func (le *LoggingExecable) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
+	start := time.Now()
+	query = le.db.mutateQuery(query)
+	row := le.UnderlyingExecable.QueryRowContext(ctx, query, args...)
+	le.db.Log.QueryTiming(ctx, "QueryRow", query, args, -1, time.Since(start), nil)
+	return row
+}
+
+func (le *LoggingExecable) Exec(query string, args ...interface{}) (sql.Result, error) {
+	return le.ExecContext(context.Background(), query, args...)
+}
+
+func (le *LoggingExecable) Query(query string, args ...interface{}) (Rows, error) {
+	return le.QueryContext(context.Background(), query, args...)
+}
+
+func (le *LoggingExecable) QueryRow(query string, args ...interface{}) *sql.Row {
+	return le.QueryRowContext(context.Background(), query, args...)
+}
+
+// loggingDB is a wrapper for LoggingExecable that allows access to BeginTx.
+//
+// While LoggingExecable has a pointer to the database and could use BeginTx, it's not technically safe since
+// the LoggingExecable could be for a transaction (where BeginTx wouldn't make sense).
+type loggingDB struct {
+	LoggingExecable
+}
+
+func (ld *loggingDB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*LoggingTxn, error) {
+	targetDB := ld.db.RawDB
+	if opts != nil && opts.ReadOnly && ld.db.ReadOnlyDB != nil {
+		targetDB = ld.db.ReadOnlyDB
+	}
+	start := time.Now()
+	tx, err := targetDB.BeginTx(ctx, opts)
+	ld.db.Log.QueryTiming(ctx, "Begin", "", nil, -1, time.Since(start), err)
+	if err != nil {
+		return nil, err
+	}
+	return &LoggingTxn{
+		LoggingExecable: LoggingExecable{UnderlyingExecable: tx, db: ld.db},
+		UnderlyingTx:    tx,
+		ctx:             ctx,
+		StartTime:       start,
+	}, nil
+}
+
+func (ld *loggingDB) Begin() (*LoggingTxn, error) {
+	return ld.BeginTx(context.Background(), nil)
+}
+
+type LoggingTxn struct {
+	LoggingExecable
+	UnderlyingTx *sql.Tx
+	ctx          context.Context
+
+	StartTime  time.Time
+	EndTime    time.Time
+	noTotalLog bool
+}
+
+func (lt *LoggingTxn) Commit() error {
+	start := time.Now()
+	err := lt.UnderlyingTx.Commit()
+	lt.EndTime = time.Now()
+	if !lt.noTotalLog {
+		lt.db.Log.QueryTiming(lt.ctx, "<Transaction>", "", nil, -1, lt.EndTime.Sub(lt.StartTime), nil)
+	}
+	lt.db.Log.QueryTiming(lt.ctx, "Commit", "", nil, -1, time.Since(start), err)
+	return err
+}
+
+func (lt *LoggingTxn) Rollback() error {
+	start := time.Now()
+	err := lt.UnderlyingTx.Rollback()
+	lt.EndTime = time.Now()
+	if !lt.noTotalLog {
+		lt.db.Log.QueryTiming(lt.ctx, "<Transaction>", "", nil, -1, lt.EndTime.Sub(lt.StartTime), nil)
+	}
+	lt.db.Log.QueryTiming(lt.ctx, "Rollback", "", nil, -1, time.Since(start), err)
+	return err
+}
+
+type LoggingRows struct {
+	ctx   context.Context
+	db    *Database
+	query string
+	args  []interface{}
+	rs    Rows
+	start time.Time
+	nrows int
+}
+
+func (lrs *LoggingRows) stopTiming() {
+	if !lrs.start.IsZero() {
+		lrs.db.Log.QueryTiming(lrs.ctx, "EndRows", lrs.query, lrs.args, lrs.nrows, time.Since(lrs.start), lrs.rs.Err())
+		lrs.start = time.Time{}
+	}
+}
+
+func (lrs *LoggingRows) Close() error {
+	err := lrs.rs.Close()
+	lrs.stopTiming()
+	return err
+}
+
+func (lrs *LoggingRows) ColumnTypes() ([]*sql.ColumnType, error) {
+	return lrs.rs.ColumnTypes()
+}
+
+func (lrs *LoggingRows) Columns() ([]string, error) {
+	return lrs.rs.Columns()
+}
+
+func (lrs *LoggingRows) Err() error {
+	return lrs.rs.Err()
+}
+
+func (lrs *LoggingRows) Next() bool {
+	hasNext := lrs.rs.Next()
+
+	if !hasNext {
+		lrs.stopTiming()
+	} else {
+		lrs.nrows++
+	}
+
+	return hasNext
+}
+
+func (lrs *LoggingRows) NextResultSet() bool {
+	hasNext := lrs.rs.NextResultSet()
+
+	if !hasNext {
+		lrs.stopTiming()
+	} else {
+		lrs.nrows++
+	}
+
+	return hasNext
+}
+
+func (lrs *LoggingRows) Scan(dest ...any) error {
+	return lrs.rs.Scan(dest...)
+}