initial commit on Federted servers

internal/db/db.go

@@ -0,0 +1,153 @@
+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()
+}