package db import ( "context" "crypto/rand" "database/sql" "encoding/hex" "encoding/json" "errors" "fmt" "sort" "strings" "time" "github.com/andrew/go-computer/internal/config" _ "github.com/jackc/pgx/v5/stdlib" ) type DB struct{ *sql.DB } type Object = map[string]any var ErrNotFound = errors.New("not found") func New(cfg *config.Config) (*DB, error) { d, err := sql.Open("pgx", cfg.DatabaseURL) if err != nil { return nil, err } d.SetMaxOpenConns(20) d.SetMaxIdleConns(5) d.SetConnMaxLifetime(time.Hour) ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() if err = d.PingContext(ctx); err != nil { d.Close() return nil, fmt.Errorf("database: %w", err) } return &DB{d}, nil } func (d *DB) Close() error { return d.DB.Close() } type contextKey struct{} // FromContext returns the request-scoped database. A missing value is a // wiring bug (route mounted without withDB) and must fail loudly here rather // than nil-panic at the first query. func FromContext(ctx context.Context) *DB { d, ok := ctx.Value(contextKey{}).(*DB) if !ok || d == nil { panic("db: no database in request context (route missing withDB middleware)") } return d } func NewContext(ctx context.Context, d *DB) context.Context { return context.WithValue(ctx, contextKey{}, d) } func NewID() string { b := make([]byte, 16) if _, err := rand.Read(b); err != nil { panic(err) } return hex.EncodeToString(b) } // Queryer is also implemented by sql.Tx, so mutations and relation updates // can be committed atomically without a second connection pool. type Queryer interface { QueryContext(context.Context, string, ...any) (*sql.Rows, error) QueryRowContext(context.Context, string, ...any) *sql.Row ExecContext(context.Context, string, ...any) (sql.Result, error) } func One(ctx context.Context, q Queryer, query string, args ...any) (Object, error) { var raw []byte err := q.QueryRowContext(ctx, query, args...).Scan(&raw) if errors.Is(err, sql.ErrNoRows) { return nil, ErrNotFound } if err != nil { return nil, err } var v Object err = json.Unmarshal(raw, &v) return v, err } func Many(ctx context.Context, q Queryer, query string, args ...any) ([]Object, error) { rows, err := q.QueryContext(ctx, query, args...) if err != nil { return nil, err } defer rows.Close() out := []Object{} for rows.Next() { var raw []byte var v Object if err = rows.Scan(&raw); err != nil { return nil, err } if err = json.Unmarshal(raw, &v); err != nil { return nil, err } out = append(out, v) } return out, rows.Err() } func Quote(s string) string { return `"` + strings.ReplaceAll(s, `"`, `""`) + `"` } // Table and column names originate in server-owned allowlists, never request // parameters. Every value is a PostgreSQL bind parameter. func Insert(ctx context.Context, q Queryer, table string, data Object) (Object, error) { keys := []string{} for k := range data { keys = append(keys, k) } sort.Strings(keys) cols, ps := []string{}, []string{} args := []any{} for i, k := range keys { cols = append(cols, Quote(k)) ps = append(ps, fmt.Sprintf("$%d", i+1)) args = append(args, data[k]) } return One(ctx, q, `INSERT INTO `+Quote(table)+` (`+strings.Join(cols, ",")+`) VALUES (`+strings.Join(ps, ",")+`) RETURNING to_jsonb(`+Quote(table)+`)`, args...) } func Update(ctx context.Context, q Queryer, table, id, ownerCol, owner string, data Object) (Object, error) { keys := []string{} for k := range data { keys = append(keys, k) } sort.Strings(keys) sets, args := []string{}, []any{} for i, k := range keys { sets = append(sets, Quote(k)+fmt.Sprintf("=$%d", i+1)) args = append(args, data[k]) } args = append(args, id, owner) where := fmt.Sprintf(`id=$%d AND %s=$%d`, len(args)-1, Quote(ownerCol), len(args)) if len(sets) == 0 { return One(ctx, q, `SELECT to_jsonb(t) FROM `+Quote(table)+` t WHERE `+where, args...) } return One(ctx, q, `UPDATE `+Quote(table)+` SET `+strings.Join(sets, ",")+` WHERE `+where+` RETURNING to_jsonb(`+Quote(table)+`)`, args...) } func Transaction(ctx context.Context, d *DB, fn func(*sql.Tx) error) error { tx, err := d.BeginTx(ctx, nil) if err != nil { return err } defer tx.Rollback() if err = fn(tx); err != nil { return err } return tx.Commit() }