Log slow API requests with server spans and client timing

internal/store/traced.go

@@ -0,0 +1,95 @@
+package store
+
+import (
+	"context"
+	"database/sql/driver"
+	"time"
+
+	"modernc.org/sqlite"
+
+	"g.fc.run/ai/byre/internal/trace"
+)
+
+type sqliteConn interface {
+	driver.Conn
+	driver.ConnBeginTx
+	driver.ConnPrepareContext
+	driver.ExecerContext
+	driver.QueryerContext
+	driver.Pinger
+	driver.SessionResetter
+	driver.Validator
+}
+
+type tracedConnector struct {
+	driver.Connector
+}
+
+func newTracedConnector(dsn string) (driver.Connector, error) {
+	c, err := sqlite.NewConnector(dsn)
+	if err != nil {
+		return nil, err
+	}
+	return tracedConnector{c}, nil
+}
+
+func (c tracedConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	conn, err := c.Connector.Connect(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return tracedConn{conn.(sqliteConn)}, nil
+}
+
+type tracedConn struct {
+	sqliteConn
+}
+
+func (c tracedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	defer trace.Start(ctx, "sql begin")()
+	return c.sqliteConn.BeginTx(ctx, opts)
+}
+
+func (c tracedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	defer trace.Start(ctx, "sql exec")()
+	return c.sqliteConn.ExecContext(ctx, query, args)
+}
+
+func (c tracedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	start := time.Now()
+	rows, err := c.sqliteConn.QueryContext(ctx, query, args)
+	if err != nil {
+		trace.Add(ctx, "sql query", time.Since(start))
+		return nil, err
+	}
+	return &tracedRows{sqliteRows: rows.(sqliteRows), ctx: ctx, took: time.Since(start)}, nil
+}
+
+type sqliteRows interface {
+	driver.Rows
+	driver.RowsColumnTypeDatabaseTypeName
+	driver.RowsColumnTypeLength
+	driver.RowsColumnTypeNullable
+	driver.RowsColumnTypePrecisionScale
+	driver.RowsColumnTypeScanType
+}
+
+type tracedRows struct {
+	sqliteRows
+	ctx  context.Context
+	took time.Duration
+}
+
+func (r *tracedRows) Next(dest []driver.Value) error {
+	start := time.Now()
+	err := r.sqliteRows.Next(dest)
+	r.took += time.Since(start)
+	return err
+}
+
+func (r *tracedRows) Close() error {
+	start := time.Now()
+	err := r.sqliteRows.Close()
+	trace.Add(r.ctx, "sql query", r.took+time.Since(start))
+	return err
+}