minstrel/internal/backup/restore.go

package backup

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var (
	// ErrNotEmpty is a restore into an instance that already has accounts or
	// a library.
	ErrNotEmpty = errors.New("backup: this Minstrel already has accounts or a library; a backup restores into an empty one")

	// ErrNewerSchema is a backup written by a newer Minstrel than this one.
	ErrNewerSchema = errors.New("backup: the backup was written by a newer version of Minstrel; upgrade this one first")

	// ErrNotABackup is contents that aren't rows a backup writes.
	ErrNotABackup = errors.New("backup: not a Minstrel backup")
)

// batchRows is how many rows one INSERT carries.
const batchRows = 500

// Restore loads a backup into an empty instance, in one transaction: either
// all of it lands or none of it does.
//
// The instance must have no accounts and no tracks. Its other rows in the
// tables the backup holds are replaced, which is what a fresh instance's
// seeded settings need. A table the backup doesn't hold, because it was
// written before that table existed, keeps what the instance has.
//
// A row's columns are those the backup and the table both have: an older
// backup's rows take the defaults of columns added since, and a column dropped
// since is left out.
func Restore(ctx context.Context, pool *pgxpool.Pool, c Contents) error {
	tx, err := pool.Begin(ctx)
	if err != nil {
		return fmt.Errorf("backup: begin: %w", err)
	}
	defer func() { _ = tx.Rollback(ctx) }()

	schema, err := schemaVersion(ctx, tx)
	if err != nil {
		return err
	}
	if c.Schema > schema {
		return ErrNewerSchema
	}
	// Taken before the emptiness check, so a first account or a scan starting
	// at the same moment waits for the restore, then finds the data.
	if _, err := tx.Exec(ctx, "LOCK TABLE users, tracks IN EXCLUSIVE MODE"); err != nil {
		return fmt.Errorf("backup: lock: %w", err)
	}
	var occupied bool
	if err := tx.QueryRow(ctx,
		"SELECT EXISTS (SELECT 1 FROM users) OR EXISTS (SELECT 1 FROM tracks)").Scan(&occupied); err != nil {
		return fmt.Errorf("backup: check the instance is empty: %w", err)
	}
	if occupied {
		return ErrNotEmpty
	}

	// Children first, so no row is left pointing at one already gone.
	for _, t := range slices.Backward(Tables) {
		if _, ok := c.Tables[t.Name]; !ok {
			continue
		}
		if _, err := tx.Exec(ctx, "DELETE FROM "+pgx.Identifier{t.Name}.Sanitize()); err != nil {
			return fmt.Errorf("backup: clear %s: %w", t.Name, err)
		}
	}
	for _, t := range Tables {
		rows, ok := c.Tables[t.Name]
		if !ok {
			continue
		}
		if err := restoreTable(ctx, tx, t.Name, rows); err != nil {
			return err
		}
	}
	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("backup: commit: %w", err)
	}
	return nil
}

type column struct {
	name     string
	identity bool
	sequence bool
}

func restoreTable(ctx context.Context, tx pgx.Tx, table string, rows [][]byte) error {
	cols, err := writableColumns(ctx, tx, table)
	if err != nil {
		return err
	}
	ident := pgx.Identifier{table}.Sanitize()
	overriding := ""
	if slices.ContainsFunc(cols, func(c column) bool { return c.identity }) {
		overriding = " OVERRIDING SYSTEM VALUE"
	}

	// Rows go in batches that share one key set; a backup's rows for a table
	// all do, since one read wrote them.
	var (
		batch []json.RawMessage
		list  string
	)
	flush := func() error {
		if len(batch) == 0 {
			return nil
		}
		payload, err := json.Marshal(batch)
		if err != nil {
			return err
		}
		sql := fmt.Sprintf("INSERT INTO %s (%s)%s SELECT %s FROM jsonb_populate_recordset(NULL::%s, $1::jsonb)",
			ident, list, overriding, list, ident)
		if _, err := tx.Exec(ctx, sql, string(payload)); err != nil {
			return fmt.Errorf("backup: restore %s: %w", table, err)
		}
		batch = batch[:0]
		return nil
	}
	for i, raw := range rows {
		var keys map[string]json.RawMessage
		if err := json.Unmarshal(raw, &keys); err != nil {
			return fmt.Errorf("%w: row %d of %s is not an object", ErrNotABackup, i+1, table)
		}
		var names []string
		for _, c := range cols {
			if _, ok := keys[c.name]; ok {
				names = append(names, pgx.Identifier{c.name}.Sanitize())
			}
		}
		if len(names) == 0 {
			continue
		}
		slices.Sort(names)
		rowList := strings.Join(names, ", ")
		if rowList != list || len(batch) == batchRows {
			if err := flush(); err != nil {
				return err
			}
			list = rowList
		}
		batch = append(batch, raw)
	}
	if err := flush(); err != nil {
		return err
	}

	// A sequence carries on after the highest restored value, so the next row
	// doesn't collide with one from the backup.
	for _, c := range cols {
		if !c.sequence {
			continue
		}
		col := pgx.Identifier{c.name}.Sanitize()
		if _, err := tx.Exec(ctx, fmt.Sprintf(
			"SELECT setval(pg_get_serial_sequence($1, $2), COALESCE((SELECT max(%s) FROM %s), 0) + 1, false)", col, ident),
			table, c.name); err != nil {
			return fmt.Errorf("backup: carry on %s's sequence: %w", table, err)
		}
	}
	return nil
}

// writableColumns returns the columns of a table an INSERT can set: all but
// generated ones, which Postgres computes.
func writableColumns(ctx context.Context, tx pgx.Tx, table string) ([]column, error) {
	rows, err := tx.Query(ctx, `
		SELECT column_name, is_identity = 'YES', pg_get_serial_sequence($1::text, column_name) IS NOT NULL
		  FROM information_schema.columns
		 WHERE table_schema = current_schema() AND table_name = $1::text AND is_generated = 'NEVER'
		 ORDER BY ordinal_position`, table)
	if err != nil {
		return nil, fmt.Errorf("backup: read %s's columns: %w", table, err)
	}
	cols, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (column, error) {
		var c column
		err := row.Scan(&c.name, &c.identity, &c.sequence)
		return c, err
	})
	if err != nil {
		return nil, fmt.Errorf("backup: read %s's columns: %w", table, err)
	}
	return cols, nil
}