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 err := warnLeftovers(ctx, opt, m); 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, ", "))
}

// warnLeftovers names the collections already in the target that the archive
// does not carry.
//
// They are named rather than dropped. A --force restore of an archive taken
// with --exclude workflow_log_lines leaves the old logs joined to restored
// runs, which the operator must know; but dropping a collection the archive
// never mentioned would delete data nobody asked to delete, and there is no
// way back from that.
func warnLeftovers(ctx context.Context, opt RestoreOptions, m Manifest) error {
	if !opt.Force {
		// Without Force the target was already proven empty.
		return nil
	}
	names, err := opt.Client.Database(opt.Database).ListCollectionNames(ctx, bson.M{})
	if err != nil {
		return fmt.Errorf("inspect target: %w", err)
	}
	inArchive := map[string]bool{}
	for _, c := range m.Collections {
		inArchive[c.Name] = true
	}
	var leftover []string
	for _, n := range names {
		if !inArchive[n] {
			leftover = append(leftover, n)
		}
	}
	if len(leftover) == 0 {
		return nil
	}
	sort.Strings(leftover)
	opt.warn("this archive does not carry %s, which already exist in %s and are left "+
		"untouched: their contents will sit alongside the restored data",
		strings.Join(leftover, ", "), opt.Database)
	return nil
}

// 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.
//
// The specs are handed to the createIndexes command exactly as the source
// server reported them, rather than reconstructed into a mongo.IndexModel from
// a hand-picked set of options. Reconstruction dropped every option nobody had
// thought to pick — partialFilterExpression above all, which this codebase
// relies on for partial unique indexes, and which replayed as a full unique
// index fails on any real database. It also lost compound key order, which is
// significant.
//
// 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 encoded []json.RawMessage
	if err := json.Unmarshal(raw, &encoded); err != nil {
		return 0, fmt.Errorf("parse index specs for %s: %w", name, err)
	}

	db := coll.Database()
	created := 0
	for _, ej := range encoded {
		spec, indexName, unique, ok, err := indexSpecFrom(ej)
		if err != nil {
			return created, fmt.Errorf("parse index specs for %s: %w", name, err)
		}
		if !ok {
			continue
		}
		cmd := bson.D{
			{Key: "createIndexes", Value: name},
			{Key: "indexes", Value: bson.A{spec}},
		}
		if err := db.RunCommand(ctx, cmd).Err(); 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
}

// droppedIndexSpecFields are the fields the server reports on an existing index
// but rejects when creating one. Everything else is passed through untouched.
var droppedIndexSpecFields = map[string]bool{"v": true, "ns": true}

// indexSpecFrom decodes one archived extended-JSON index specification into an
// ordered bson.D suitable for createIndexes.
//
// The _id_ index is skipped: MongoDB creates it itself and refuses an explicit
// attempt to create it.
func indexSpecFrom(ej []byte) (bson.D, string, bool, bool, error) {
	var d bson.D
	if err := bson.UnmarshalExtJSON(ej, false, &d); err != nil {
		return nil, "", false, false, err
	}

	out := make(bson.D, 0, len(d))
	var name string
	unique := false
	hasKey := false
	for _, e := range d {
		switch e.Key {
		case "name":
			name, _ = e.Value.(string)
		case "unique":
			if u, ok := e.Value.(bool); ok {
				unique = u
			}
		case "key":
			hasKey = true
		}
		if droppedIndexSpecFields[e.Key] {
			continue
		}
		out = append(out, e)
	}
	if name == "_id_" || !hasKey {
		return nil, name, false, false, nil
	}
	return out, name, unique, true, nil
}