96 lines
2.1 KiB
Go
96 lines
2.1 KiB
Go
|
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
|
||
|
|
}
|