vantage-shared/backup/restore.go

package backup

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"sort"
	"strings"

	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

// BatchSize is how many documents are inserted per bulk write.
const BatchSize = 1000

// ErrTargetNotEmpty is returned when the target database already holds data and
// Force was not set.
var ErrTargetNotEmpty = errors.New("target database is not empty")

// ErrKeyMismatch is returned when the archive's key fingerprint does not match
// the key supplied.
var ErrKeyMismatch = errors.New("KEY_ENCRYPTION_KEY does not match the archive")

// ErrIndexBuild is returned when a unique index in the archive cannot be built
// on the restored data.
var ErrIndexBuild = errors.New("index could not be built on the restored data")

// RestoredCollection is what one collection's restore produced.
type RestoredCollection struct {
	Name      string
	Documents int64
	Indexes   int
}

// RestoreResult is the summary a caller prints.
type RestoreResult struct {
	Collections []RestoredCollection
}

// RestoreOptions configures one restore.
type RestoreOptions struct {
	Client   *mongo.Client
	Database string
	Archive  *Reader

	// Force drops each collection in the archive before loading it. Without it
	// a non-empty target is refused.
	Force bool

	KeyHex string

	// IgnoreKeyMismatch proceeds past a fingerprint mismatch, having first
	// warned which collections will hold unreadable ciphertext afterwards.
	IgnoreKeyMismatch bool

	// Warn receives operator-facing warnings. A nil Warn discards them.
	Warn func(string)
}

func (o RestoreOptions) warn(format string, args ...any) {
	if o.Warn != nil {
		o.Warn(fmt.Sprintf(format, args...))
	}
}

// Restore loads an archive into a database.
//
// The order is fixed and every check that can refuse does so before the first
// write: format, checksums (done by Open), key policy, then target inspection.
// A restore that has begun writing and then fails leaves a partial database
// which the next run refuses to touch, which is correct — the alternative is a
// silent merge, and merging two control planes reconciles nothing.
func Restore(ctx context.Context, opt RestoreOptions) (RestoreResult, error) {
	m := opt.Archive.Manifest()

	if err := checkKey(m, opt); err != nil {
		return RestoreResult{}, err
	}
	if err := checkTarget(ctx, opt); err != nil {
		return RestoreResult{}, err
	}
	if len(m.Excluded) > 0 {
		opt.warn("this archive excluded %s; those collections will be empty after the restore",
			strings.Join(m.Excluded, ", "))
	}
	opt.warn("Redis is not restored. Sessions are the only state it holds, so everyone signs in again.")

	res := RestoreResult{}
	if err := warnVersionGap(ctx, opt, m); err != nil {
		return res, err
	}

	for _, entry := range m.Collections {
		rc, err := restoreCollection(ctx, opt, entry)
		if err != nil {
			return res, err
		}
		res.Collections = append(res.Collections, rc)
	}
	return res, nil
}

// checkKey applies the fingerprint policy.
func checkKey(m Manifest, opt RestoreOptions) error {
	if m.KeyFingerprint == nil {
		opt.warn("this archive carries no key fingerprint, so nothing here proves your " +
			"KEY_ENCRYPTION_KEY opens its ciphertext")
		return nil
	}
	if opt.KeyHex == "" {
		return fmt.Errorf("%w: the archive records a key fingerprint, so a key is required "+
			"(pass --ignore-key-mismatch only if you accept unreadable secrets)", ErrNoKey)
	}
	got, err := FingerprintHex(opt.KeyHex)
	if err != nil {
		return err
	}
	if got == *m.KeyFingerprint {
		return nil
	}
	if !opt.IgnoreKeyMismatch {
		return fmt.Errorf("%w: archive fingerprint %s, your key fingerprints as %s",
			ErrKeyMismatch, *m.KeyFingerprint, got)
	}
	opt.warn("proceeding past a key mismatch: ciphertext in %s will be permanently unreadable",
		strings.Join(CiphertextCollections(), ", "))
	return nil
}

// checkTarget refuses a non-empty database unless Force was set.
func checkTarget(ctx context.Context, opt RestoreOptions) error {
	db := opt.Client.Database(opt.Database)
	names, err := db.ListCollectionNames(ctx, bson.M{})
	if err != nil {
		return fmt.Errorf("inspect target: %w", err)
	}
	if len(names) == 0 || opt.Force {
		return nil
	}
	sort.Strings(names)

	var found []string
	for _, n := range names {
		count, err := db.Collection(n).CountDocuments(ctx, bson.M{})
		if err != nil {
			return fmt.Errorf("count %s: %w", n, err)
		}
		found = append(found, fmt.Sprintf("%s (%d)", n, count))
	}
	return fmt.Errorf("%w: %s holds %s", ErrTargetNotEmpty, opt.Database, strings.Join(found, ", "))
}

// warnVersionGap reports a major version difference between the server that
// produced the archive and the one receiving it. It warns rather than refuses:
// restoring across a major version is a normal part of an upgrade, and a tool
// that refused would be blocking the migration it exists to make safe.
func warnVersionGap(ctx context.Context, opt RestoreOptions, m Manifest) error {
	if m.MongoServerVersion == "" {
		return nil
	}
	target, err := mongoServerVersion(ctx, opt.Client)
	if err != nil {
		return err
	}
	if majorOf(m.MongoServerVersion) != majorOf(target) {
		opt.warn("this archive came from MongoDB %s and you are restoring onto %s",
			m.MongoServerVersion, target)
	}
	return nil
}

func majorOf(version string) string {
	if i := strings.IndexByte(version, '.'); i >= 0 {
		return version[:i]
	}
	return version
}

func restoreCollection(ctx context.Context, opt RestoreOptions, entry CollectionEntry) (RestoredCollection, error) {
	coll := opt.Client.Database(opt.Database).Collection(entry.Name)
	if opt.Force {
		if err := coll.Drop(ctx); err != nil {
			return RestoredCollection{}, fmt.Errorf("drop %s: %w", entry.Name, err)
		}
	}

	written, err := insertDocuments(ctx, opt, coll, entry)
	if err != nil {
		return RestoredCollection{}, err
	}
	indexes, err := replayIndexes(ctx, opt, coll, entry.Name)
	if err != nil {
		return RestoredCollection{}, err
	}
	return RestoredCollection{Name: entry.Name, Documents: written, Indexes: indexes}, nil
}

func insertDocuments(ctx context.Context, opt RestoreOptions, coll *mongo.Collection, entry CollectionEntry) (int64, error) {
	rc, err := opt.Archive.OpenCollection(entry.Name)
	if err != nil {
		return 0, fmt.Errorf("open %s in archive: %w", entry.Name, err)
	}
	defer rc.Close()

	raw, err := io.ReadAll(rc)
	if err != nil {
		return 0, fmt.Errorf("read %s: %w", entry.Name, err)
	}

	var written int64
	batch := make([]any, 0, BatchSize)
	flush := func() error {
		if len(batch) == 0 {
			return nil
		}
		if _, err := coll.InsertMany(ctx, batch, options.InsertMany().SetOrdered(false)); err != nil {
			return fmt.Errorf("insert into %s: %w", entry.Name, err)
		}
		written += int64(len(batch))
		batch = batch[:0]
		return nil
	}

	for len(raw) > 0 {
		doc, rest, err := splitBSON(raw)
		if err != nil {
			return 0, fmt.Errorf("%s: %w", entry.Name, err)
		}
		batch = append(batch, doc)
		raw = rest
		if len(batch) == BatchSize {
			if err := flush(); err != nil {
				return 0, err
			}
		}
	}
	if err := flush(); err != nil {
		return 0, err
	}
	return written, nil
}

// splitBSON peels one document off the front of a concatenated BSON stream. A
// BSON document declares its own length in its first four bytes.
func splitBSON(raw []byte) (bson.Raw, []byte, error) {
	if len(raw) < 4 {
		return nil, nil, fmt.Errorf("truncated BSON: %d trailing bytes", len(raw))
	}
	n := int(int32(raw[0]) | int32(raw[1])<<8 | int32(raw[2])<<16 | int32(raw[3])<<24)
	if n < 5 || n > len(raw) {
		return nil, nil, fmt.Errorf("BSON document declares length %d with %d bytes remaining", n, len(raw))
	}
	return bson.Raw(raw[:n]), raw[n:], nil
}

// replayIndexes recreates the archived indexes.
//
// A unique index that will not build means the restored data violates it, and
// the unique indexes here — (instance_id, email), instance slug, settings
// instance, the ESO token hash — are tenant-isolation properties rather than
// optimisations. That aborts. A non-unique index failing is a performance
// problem and warns.
func replayIndexes(ctx context.Context, opt RestoreOptions, coll *mongo.Collection, name string) (int, error) {
	raw, err := opt.Archive.IndexesJSON(name)
	if err != nil {
		return 0, fmt.Errorf("read index specs for %s: %w", name, err)
	}
	if len(raw) == 0 {
		return 0, nil
	}
	var specs []map[string]any
	if err := json.Unmarshal(raw, &specs); err != nil {
		return 0, fmt.Errorf("parse index specs for %s: %w", name, err)
	}

	created := 0
	for _, spec := range specs {
		model, indexName, unique, ok := indexModelFrom(spec)
		if !ok {
			continue
		}
		if _, err := coll.Indexes().CreateOne(ctx, model); err != nil {
			if unique {
				return created, fmt.Errorf("%w: %s on %s: %v", ErrIndexBuild, indexName, name, err)
			}
			opt.warn("index %s on %s was not created: %v", indexName, name, err)
			continue
		}
		created++
	}
	return created, nil
}

// indexModelFrom converts one archived index specification into a model.
// The _id_ index is skipped: MongoDB creates it itself and refuses an explicit
// attempt to create it.
func indexModelFrom(spec map[string]any) (mongo.IndexModel, string, bool, bool) {
	name, _ := spec["name"].(string)
	if name == "_id_" {
		return mongo.IndexModel{}, name, false, false
	}
	keys, ok := spec["key"].(map[string]any)
	if !ok || len(keys) == 0 {
		return mongo.IndexModel{}, name, false, false
	}

	// JSON objects do not preserve order but compound index key order is
	// significant, so the field order recorded by the server is recovered from
	// the spec's own ordering where available and sorted otherwise. bson.M
	// round-trips through json as a map; the archive therefore stores the key
	// document and this reconstructs a deterministic bson.D from it.
	fields := make([]string, 0, len(keys))
	for k := range keys {
		fields = append(fields, k)
	}
	sort.Strings(fields)
	d := make(bson.D, 0, len(fields))
	for _, f := range fields {
		d = append(d, bson.E{Key: f, Value: keys[f]})
	}

	opts := options.Index().SetName(name)
	unique := false
	if u, ok := spec["unique"].(bool); ok && u {
		unique = true
		opts = opts.SetUnique(true)
	}
	if s, ok := spec["sparse"].(bool); ok && s {
		opts = opts.SetSparse(true)
	}
	if e, ok := spec["expireAfterSeconds"].(float64); ok {
		opts = opts.SetExpireAfterSeconds(int32(e))
	}
	return mongo.IndexModel{Keys: d, Options: opts}, name, unique, true
}