fix: replay index specs verbatim instead of reconstructing them

2026-09-07 14:45:36 +00:00
parent d328d3aaca
commit 2e4bc687d4
3 changed files with 214 additions and 55 deletions
@@ -147,11 +147,25 @@ func dumpIndexes(ctx context.Context, w *Writer, db *mongo.Database, name string
 	}
 	defer cur.Close(ctx)

-	var specs []bson.M
+	// The specs are read as raw BSON and re-encoded as extended JSON, one
+	// element per index, so key order and every option the server reported —
+	// partialFilterExpression, collation, weights and the rest — survive
+	// verbatim. Decoding into bson.M would lose compound key order, and
+	// reconstructing an index from a hand-picked set of options would drop
+	// whatever was not picked.
+	var specs []bson.Raw
 	if err := cur.All(ctx, &specs); err != nil {
 		return fmt.Errorf("read indexes on %s: %w", name, err)
 	}
-	raw, err := json.Marshal(specs)
+	encoded := make([]json.RawMessage, 0, len(specs))
+	for _, spec := range specs {
+		ej, err := bson.MarshalExtJSON(spec, false, false)
+		if err != nil {
+			return fmt.Errorf("encode indexes on %s: %w", name, err)
+		}
+		encoded = append(encoded, ej)
+	}
+	raw, err := json.Marshal(encoded)
 	if err != nil {
 		return fmt.Errorf("encode indexes on %s: %w", name, err)
 	}
@@ -83,6 +83,9 @@ func Restore(ctx context.Context, opt RestoreOptions) (RestoreResult, error) {
 	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, ", "))
@@ -154,6 +157,43 @@ func checkTarget(ctx context.Context, opt RestoreOptions) error {
 	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
@@ -259,6 +299,14 @@ func splitBSON(raw []byte) (bson.Raw, []byte, error) {

 // 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
@@ -272,18 +320,26 @@ func replayIndexes(ctx context.Context, opt RestoreOptions, coll *mongo.Collecti
 	if len(raw) == 0 {
 		return 0, nil
 	}
-	var specs []map[string]any
-	if err := json.Unmarshal(raw, &specs); err != 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 _, spec := range specs {
-		model, indexName, unique, ok := indexModelFrom(spec)
+	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
 		}
-		if _, err := coll.Indexes().CreateOne(ctx, model); err != nil {
+		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)
 			}
@@ -295,45 +351,43 @@ func replayIndexes(ctx context.Context, opt RestoreOptions, coll *mongo.Collecti
 	return created, nil
 }

-// indexModelFrom converts one archived index specification into a model.
+// 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 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
+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
 	}

-	// 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)
+	out := make(bson.D, 0, len(d))
+	var name string
 	unique := false
-	if u, ok := spec["unique"].(bool); ok && u {
-		unique = true
-		opts = opts.SetUnique(true)
+	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 s, ok := spec["sparse"].(bool); ok && s {
-		opts = opts.SetSparse(true)
+	if name == "_id_" || !hasKey {
+		return nil, name, false, false, nil
 	}
-	if e, ok := spec["expireAfterSeconds"].(float64); ok {
-		opts = opts.SetExpireAfterSeconds(int32(e))
-	}
-	return mongo.IndexModel{Keys: d, Options: opts}, name, unique, true
+	return out, name, unique, true, nil
 }
@@ -11,6 +11,7 @@ import (

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

 // archiveOf seeds a database, dumps it, and returns an opened Reader.
@@ -320,11 +321,12 @@ func TestRestoreSameVersionDoesNotWarnAboutIt(t *testing.T) {
 	}
 }

-func TestCompoundIndexKeyOrderIsAlphabetical(t *testing.T) {
-	model, name, unique, ok := indexModelFrom(map[string]any{
-		"name": "b_1_a_1",
-		"key":  map[string]any{"b": float64(1), "a": float64(1)},
-	})
+func TestCompoundIndexKeyOrderIsPreserved(t *testing.T) {
+	spec, name, unique, ok, err := indexSpecFrom([]byte(
+		`{"v":2,"key":{"b":1,"a":1},"name":"b_1_a_1","ns":"db.c"}`))
+	if err != nil {
+		t.Fatalf("indexSpecFrom: %v", err)
+	}
 	if !ok {
 		t.Fatal("spec rejected")
 	}
@@ -334,15 +336,104 @@ func TestCompoundIndexKeyOrderIsAlphabetical(t *testing.T) {
 	if name != "b_1_a_1" {
 		t.Fatalf("name %q", name)
 	}
-	keys, isD := model.Keys.(bson.D)
-	if !isD {
-		t.Fatalf("keys are %T, want bson.D", model.Keys)
+
+	var keys bson.D
+	for _, e := range spec {
+		switch e.Key {
+		case "v", "ns":
+			t.Fatalf("%q must be stripped before createIndexes, got %v", e.Key, spec)
+		case "key":
+			d, isD := e.Value.(bson.D)
+			if !isD {
+				t.Fatalf("key is %T, want bson.D", e.Value)
+			}
+			keys = d
+		}
 	}
-	// Documents current behaviour: field order is alphabetical, not the order
-	// the server reported. Storing the key document as raw BSON in the archive
-	// instead of JSON would fix this and is the change to make if compound
-	// index order ever matters here.
-	if keys[0].Key != "a" || keys[1].Key != "b" {
-		t.Fatalf("got %v", keys)
+	// Compound index key order is significant, so it is carried through
+	// verbatim rather than reconstructed from an unordered map.
+	if len(keys) != 2 || keys[0].Key != "b" || keys[1].Key != "a" {
+		t.Fatalf("key order not preserved, got %v", keys)
 	}
 }
+
+func TestIdIndexIsSkipped(t *testing.T) {
+	_, name, _, ok, err := indexSpecFrom([]byte(`{"v":2,"key":{"_id":1},"name":"_id_"}`))
+	if err != nil {
+		t.Fatalf("indexSpecFrom: %v", err)
+	}
+	if ok {
+		t.Fatal("_id_ must be skipped; MongoDB creates it itself")
+	}
+	if name != "_id_" {
+		t.Fatalf("name %q", name)
+	}
+}
+
+// TestRestoreReplaysPartialUniqueIndex is the regression guard for the defect
+// that made a restore abort on any real database: a partial unique index —
+// this codebase has them on workflow_steps and settings — replayed as a full
+// unique index hits duplicate keys, and a failing unique index is fatal.
+func TestRestoreReplaysPartialUniqueIndex(t *testing.T) {
+	client, _ := testDB(t)
+	ctx := context.Background()
+	_, srcDB := testDB(t)
+	seed(t, client, srcDB)
+
+	coll := client.Database(srcDB).Collection("workflow_steps")
+	docs := []any{
+		bson.M{"instance_id": "i1", "slug": "same", "source": "default"},
+		bson.M{"instance_id": "i1", "slug": "same", "source": "custom"},
+		bson.M{"instance_id": "i1", "slug": "same", "source": "custom"},
+	}
+	if _, err := coll.InsertMany(ctx, docs); err != nil {
+		t.Fatalf("insert: %v", err)
+	}
+	if _, err := coll.Indexes().CreateOne(ctx, mongo.IndexModel{
+		Keys: bson.D{{Key: "instance_id", Value: 1}, {Key: "slug", Value: 1}},
+		Options: options.Index().SetName("default_step_slug").SetUnique(true).
+			SetPartialFilterExpression(bson.M{"source": "default"}),
+	}); err != nil {
+		t.Fatalf("create partial index: %v", err)
+	}
+
+	path, _ := dumpToFile(t, DumpOptions{Client: client, Database: srcDB, KeyHex: validKeyHex})
+	archive, err := Open(path)
+	if err != nil {
+		t.Fatalf("Open: %v", err)
+	}
+	defer archive.Close()
+
+	_, target := testDB(t)
+	if _, err := Restore(ctx, RestoreOptions{
+		Client: client, Database: target, Archive: archive, KeyHex: validKeyHex,
+	}); err != nil {
+		t.Fatalf("Restore: %v", err)
+	}
+
+	cur, err := client.Database(target).Collection("workflow_steps").Indexes().List(ctx)
+	if err != nil {
+		t.Fatalf("list indexes: %v", err)
+	}
+	var specs []bson.M
+	if err := cur.All(ctx, &specs); err != nil {
+		t.Fatalf("read indexes: %v", err)
+	}
+	for _, s := range specs {
+		if s["name"] != "default_step_slug" {
+			continue
+		}
+		if s["unique"] != true {
+			t.Fatalf("index lost its uniqueness: %v", s)
+		}
+		if s["partialFilterExpression"] == nil {
+			t.Fatalf("partialFilterExpression was dropped: %v", s)
+		}
+		keys, isD := s["key"].(bson.D)
+		if isD && (len(keys) != 2 || keys[0].Key != "instance_id" || keys[1].Key != "slug") {
+			t.Fatalf("compound key order not preserved: %v", keys)
+		}
+		return
+	}
+	t.Fatalf("partial unique index not replayed; got %v", specs)
+}