Correctly measure duration of Query

db/database.go

@@ -85,7 +85,7 @@ func (d *PgxDatabase) performAndReport(fields logging.Fields, fn func() (interfa
 	ret, err := fn()
 	duration := time.Since(start)
 	if duration > 100*time.Millisecond {
-		logger := d.getLogger().WithFields(fields).WithField("duration", duration)
+		logger := d.getLogger().WithFields(fields).WithField("took", duration)
 		if err != nil {
 			logger = logger.WithError(err)
 		}
@@ -117,15 +117,17 @@ func (d *PgxDatabase) GetPrimitive(dest interface{}, query string, args ...inter
 	return nil
 }
 func (d *PgxDatabase) Query(query string, args ...interface{}) (rows pgx.Rows, err error) {
-	ret, err := d.performAndReport(logging.Fields{
+	l := d.getLogger().WithFields(logging.Fields{
 		"type":  "start query",
 		"query": query,
 		"args":  args,
-	}, func() (interface{}, error) { return d.db.Query(d.getContext(), query, args...) })
+	})
+	start := time.Now()
+	ret, err := d.db.Query(d.getContext(), query, args...)
 	if ret == nil {
 		return nil, err
 	}
-	return ret.(pgx.Rows), err
+	return Logged(ret, start, l), err
 }
 
 func (d *PgxDatabase) Select(results interface{}, query string, args ...interface{}) error {

db/logged_rows.go

@@ -0,0 +1,74 @@
+package db
+
+import (
+	"github.com/treeverse/lakefs/logging"
+
+	"time"
+
+	"github.com/jackc/pgconn"
+	"github.com/jackc/pgproto3/v2"
+	"github.com/jackc/pgx/v4"
+)
+
+// LoggedRows is a pgx.Rows that wraps and traces another pgx.Rows.
+type LoggedRows struct {
+	pgx.Rows
+	Start  time.Time
+	Closed bool
+	l      logging.Logger
+}
+
+var _ pgx.Rows = &LoggedRows{}
+
+// Logged returns a pgx.Rows that will forward calls to r and logs their durations.
+func Logged(r pgx.Rows, start time.Time, l logging.Logger) *LoggedRows {
+	return &LoggedRows{
+		Rows:  r,
+		Start: start,
+		l:     l,
+	}
+}
+
+func (lr *LoggedRows) logDuration() {
+	if !lr.Closed {
+		lr.l.WithField("duration", time.Since(lr.Start)).Info("rows done")
+	}
+}
+
+func (lr *LoggedRows) Close() {
+	lr.Rows.Close()
+	lr.logDuration()
+}
+
+func (lr *LoggedRows) Next() bool {
+	ret := lr.Rows.Next()
+	if !ret {
+		// Underlying Rows closed itself, lr.Close() was not called
+		lr.logDuration()
+	}
+	return ret
+}
+
+func (lr *LoggedRows) Err() error {
+	return lr.Rows.Err()
+}
+
+func (lr *LoggedRows) CommandTag() pgconn.CommandTag {
+	return lr.Rows.CommandTag()
+}
+
+func (lr *LoggedRows) FieldDescriptions() []pgproto3.FieldDescription {
+	return lr.Rows.FieldDescriptions()
+}
+
+func (lr *LoggedRows) Scan(dest ...interface{}) error {
+	return lr.Rows.Scan(dest...)
+}
+
+func (lr *LoggedRows) Values() ([]interface{}, error) {
+	return lr.Rows.Values()
+}
+
+func (lr *LoggedRows) RawValues() [][]byte {
+	return lr.Rows.RawValues()
+}

db/tx.go

@@ -42,14 +42,13 @@ func (d *dbTx) Query(query string, args ...interface{}) (pgx.Rows, error) {
 		"type":  "query",
 		"args":  args,
 		"query": queryToString(query),
-		"took":  time.Since(start),
 	})
 	if err != nil {
 		log.WithError(err).Error("SQL query failed with error")
 		return nil, err
 	}
 	log.Trace("SQL query started successfully")
-	return rows, nil
+	return Logged(rows, start, log), nil
 }
 
 func Select(d Tx, results interface{}, query string, args ...interface{}) error {

go.mod

@@ -38,7 +38,7 @@ require (
 	github.com/influxdata/tdigest v0.0.1 // indirect
 	github.com/jackc/pgconn v1.6.4
 	github.com/jackc/pgerrcode v0.0.0-20190803225404-afa3381909a6
-	github.com/jackc/pgproto3/v2 v2.0.4 // indirect
+	github.com/jackc/pgproto3/v2 v2.0.4
 	github.com/jackc/pgtype v1.4.2
 	github.com/jackc/pgx/v4 v4.8.1
 	github.com/jamiealquiza/tachymeter v2.0.0+incompatible