internal/store/store.go

32a5f76e5b270097b63a5bba9a43557cf50ad63d
gitbay/internal/store/store.go history · blame · raw

290 lines · 8450 bytes

  1// Package store owns SQLite access and schema migrations.
  2package store
  3
  4import (
  5	"context"
  6	"database/sql"
  7	"embed"
  8	"errors"
  9	"fmt"
 10	"io/fs"
 11	"os"
 12	"sort"
 13	"strconv"
 14	"strings"
 15	"sync"
 16
 17	"modernc.org/sqlite"
 18)
 19
 20//go:embed migrations/*.sql
 21var migrationFS embed.FS
 22
 23type Store struct {
 24	DB *sql.DB
 25
 26	// logWait holds one channel per build someone is following, closed
 27	// by the next change to that build's row (BuildLogWait).
 28	logMu   sync.Mutex
 29	logWait map[int64]chan struct{}
 30
 31	// onRevoke runs after each key revocation this process commits.
 32	revokeMu sync.Mutex
 33	onRevoke []func(Revoked)
 34}
 35
 36// Open opens (creating if needed) the database at path with WAL mode and
 37// foreign keys enforced. Use ":memory:" in tests.
 38//
 39// _txlock=immediate is what serialises writers. Every transaction in this
 40// package writes, and a deferred one takes the write lock only when it
 41// reaches its first write — by which point another writer may hold it.
 42// SQLite answers that with SQLITE_BUSY and does not invoke the busy
 43// handler, because waiting would deadlock two transactions each holding a
 44// read lock the other needs; busy_timeout cannot help. Measured with
 45// eight concurrent read-then-write transactions, 44% of them failed.
 46// Beginning immediate takes the write lock up front, where busy_timeout
 47// does apply, so a second writer waits its turn: the same load runs with
 48// no failures, and readers, which WAL keeps out of the way, are
 49// unaffected (#121).
 50func Open(path string) (*Store, error) {
 51	dsn := path + "?_txlock=immediate&_pragma=journal_mode(WAL)&_pragma=foreign_keys(ON)&_pragma=busy_timeout(5000)"
 52	if path == ":memory:" {
 53		dsn = ":memory:?_txlock=immediate&_pragma=foreign_keys(ON)"
 54	}
 55	db, err := sql.Open("sqlite", dsn)
 56	if err != nil {
 57		return nil, err
 58	}
 59	if err := db.Ping(); err != nil {
 60		db.Close()
 61		return nil, err
 62	}
 63	// SQLite creates the file 0666&~umask, so it lands 0644 by default. The
 64	// directory above it is the real boundary, but the file holds token
 65	// hashes, addresses and private repo names and has no business being
 66	// world-readable on its own.
 67	if path != ":memory:" {
 68		if err := os.Chmod(path, 0o640); err != nil && !errors.Is(err, fs.ErrNotExist) {
 69			db.Close()
 70			return nil, err
 71		}
 72	}
 73	return &Store{DB: db}, nil
 74}
 75
 76func (s *Store) Close() error { return s.DB.Close() }
 77
 78type migration struct {
 79	version int
 80	name    string
 81	up      string
 82	down    string
 83	// upFKOff and downFKOff are true when the up/down script's first line
 84	// is the directive "-- foreign_keys: off".
 85	upFKOff   bool
 86	downFKOff bool
 87}
 88
 89// fkOffDirective, as the first line of a migration script, opts that
 90// direction out of foreign-key enforcement for its step.
 91const fkOffDirective = "-- foreign_keys: off"
 92
 93func loadMigrations() ([]migration, error) {
 94	entries, err := fs.ReadDir(migrationFS, "migrations")
 95	if err != nil {
 96		return nil, err
 97	}
 98	byVersion := map[int]*migration{}
 99	for _, e := range entries {
100		name := e.Name()
101		// <version>_<name>.<up|down>.sql
102		base, ok := strings.CutSuffix(name, ".sql")
103		if !ok {
104			return nil, fmt.Errorf("migration %q: not .sql", name)
105		}
106		var dir string
107		if b, ok := strings.CutSuffix(base, ".up"); ok {
108			base, dir = b, "up"
109		} else if b, ok := strings.CutSuffix(base, ".down"); ok {
110			base, dir = b, "down"
111		} else {
112			return nil, fmt.Errorf("migration %q: missing .up/.down", name)
113		}
114		verStr, rest, ok := strings.Cut(base, "_")
115		if !ok {
116			return nil, fmt.Errorf("migration %q: missing version prefix", name)
117		}
118		ver, err := strconv.Atoi(verStr)
119		if err != nil {
120			return nil, fmt.Errorf("migration %q: bad version: %w", name, err)
121		}
122		m := byVersion[ver]
123		if m == nil {
124			m = &migration{version: ver, name: rest}
125			byVersion[ver] = m
126		}
127		sqlBytes, err := migrationFS.ReadFile("migrations/" + name)
128		if err != nil {
129			return nil, err
130		}
131		text := string(sqlBytes)
132		firstLine, _, _ := strings.Cut(text, "\n")
133		fkOff := strings.TrimSpace(firstLine) == fkOffDirective
134		if dir == "up" {
135			m.up = text
136			m.upFKOff = fkOff
137		} else {
138			m.down = text
139			m.downFKOff = fkOff
140		}
141	}
142	var ms []migration
143	for _, m := range byVersion {
144		if m.up == "" || m.down == "" {
145			return nil, fmt.Errorf("migration %d %q: missing up or down file", m.version, m.name)
146		}
147		ms = append(ms, *m)
148	}
149	sort.Slice(ms, func(i, j int) bool { return ms[i].version < ms[j].version })
150	for i, m := range ms {
151		if m.version != i+1 {
152			return nil, fmt.Errorf("migration versions not contiguous at %d", m.version)
153		}
154	}
155	return ms, nil
156}
157
158// Version returns the current schema version (0 = empty database).
159func (s *Store) Version() (int, error) {
160	var v int
161	err := s.DB.QueryRow("PRAGMA user_version").Scan(&v)
162	return v, err
163}
164
165// MigrateUp applies all pending migrations.
166func (s *Store) MigrateUp() error { return s.migrateTo(-1) }
167
168// MigrateTo migrates up or down to the given version. 0 empties the schema.
169func (s *Store) MigrateTo(target int) error { return s.migrateTo(target) }
170
171func (s *Store) migrateTo(target int) error {
172	ms, err := loadMigrations()
173	if err != nil {
174		return err
175	}
176	if target < 0 {
177		target = len(ms)
178	}
179	if target > len(ms) {
180		return fmt.Errorf("no such schema version %d (max %d)", target, len(ms))
181	}
182	cur, err := s.Version()
183	if err != nil {
184		return err
185	}
186	step := func(sqlText string, newVersion int, fkOff bool) (retErr error) {
187		if !fkOff {
188			tx, err := s.DB.Begin()
189			if err != nil {
190				return err
191			}
192			defer tx.Rollback()
193			if _, err := tx.Exec(sqlText); err != nil {
194				return err
195			}
196			if _, err := tx.Exec(fmt.Sprintf("PRAGMA user_version = %d", newVersion)); err != nil {
197				return err
198			}
199			return tx.Commit()
200		}
201
202		// A script whose first line is "-- foreign_keys: off" rebuilds a
203		// table that other tables reference (labels, milestones): with
204		// foreign keys on, the rebuild-by-rename loses the children's
205		// rows. PRAGMA foreign_keys is a no-op inside a transaction, and
206		// the pool gives no guarantee that a pragma set on one connection
207		// is seen by the connection Begin() draws next, so the whole step
208		// — pragma off, transaction, pragma on, foreign_key_check — runs
209		// on a single pinned connection.
210		ctx := context.Background()
211		conn, err := s.DB.Conn(ctx)
212		if err != nil {
213			return err
214		}
215		defer conn.Close()
216		if _, err := conn.ExecContext(ctx, "PRAGMA foreign_keys = OFF"); err != nil {
217			return err
218		}
219		// The connection goes back to the pool when this returns, so every
220		// path out of here has to put foreign keys back on first.
221		restoreFK := func() error {
222			_, err := conn.ExecContext(ctx, "PRAGMA foreign_keys = ON")
223			return err
224		}
225		defer func() {
226			if err := restoreFK(); err != nil && retErr == nil {
227				retErr = err
228			}
229		}()
230		tx, err := conn.BeginTx(ctx, nil)
231		if err != nil {
232			return err
233		}
234		defer tx.Rollback()
235		if _, err := tx.Exec(sqlText); err != nil {
236			return err
237		}
238		if _, err := tx.Exec(fmt.Sprintf("PRAGMA user_version = %d", newVersion)); err != nil {
239			return err
240		}
241		if err := tx.Commit(); err != nil {
242			return err
243		}
244		if err := restoreFK(); err != nil {
245			return err
246		}
247		rows, err := conn.QueryContext(ctx, "PRAGMA foreign_key_check")
248		if err != nil {
249			return err
250		}
251		defer rows.Close()
252		if rows.Next() {
253			var table string
254			var rowid sql.NullInt64
255			var referredTable string
256			var fkid int
257			if err := rows.Scan(&table, &rowid, &referredTable, &fkid); err != nil {
258				return err
259			}
260			return fmt.Errorf("foreign_key_check failed after migration: %s", table)
261		}
262		return rows.Err()
263	}
264	for cur < target {
265		m := ms[cur]
266		if err := step(m.up, m.version, m.upFKOff); err != nil {
267			return fmt.Errorf("migration %d up: %w", m.version, err)
268		}
269		cur = m.version
270	}
271	for cur > target {
272		m := ms[cur-1]
273		if err := step(m.down, m.version-1, m.downFKOff); err != nil {
274			return fmt.Errorf("migration %d down: %w", m.version, err)
275		}
276		cur = m.version - 1
277	}
278	return nil
279}
280
281// IsInternal reports whether err is the database or the I/O beneath it
282// failing, as opposed to a sentinel or a message about the caller's
283// input. Callers map it to a failure exit rather than a usage error.
284func IsInternal(err error) bool {
285	var sqlErr *sqlite.Error
286	var pathErr *fs.PathError
287	return errors.As(err, &sqlErr) || errors.As(err, &pathErr) ||
288		errors.Is(err, sql.ErrTxDone) || errors.Is(err, sql.ErrConnDone) ||
289		errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled)
290}