Merge pull request #131 from FGasper/felipe_simpler_doc_key

Author: FGasper

dockey/agg.go

@@ -4,10 +4,7 @@
 package dockey
 
 import (
-	"maps"
-	"slices"
-	"strconv"
-
+	"github.com/10gen/migration-verifier/mslices"
 	"github.com/samber/lo"
 	"go.mongodb.org/mongo-driver/bson"
 )
@@ -21,69 +18,12 @@ import (
 func ExtractTrueDocKeyAgg(fieldNames []string, docExpr string) bson.D {
 	assertFieldNameUniqueness(fieldNames)
 
-	var docKeyNumKeys bson.D
-	numToKeyLookup := map[string]string{}
-
-	for n, name := range fieldNames {
-		var valExpr = docExpr + "." + name
-
-		// Aggregation forbids direct creation of an object with dotted keys.
-		// So here we create an object with numeric keys, then below we’ll
-		// map the numeric keys back to the real ones.
-
-		nStr := strconv.Itoa(n)
-		docKeyNumKeys = append(docKeyNumKeys, bson.E{nStr, valExpr})
-		numToKeyLookup[nStr] = name
-	}
-
-	// Now convert the numeric keys back to the real ones.
-	return mapObjectKeysAgg(docKeyNumKeys, numToKeyLookup)
-}
-
-// Potentially reusable:
-func mapObjectKeysAgg(expr any, mapping map[string]string) bson.D {
-	// We would ideally pass mapping into the aggregation and $getField
-	// to get the mapped key, but pre-v8 server versions required $getField’s
-	// field parameter to be a constant. (And pre-v5 didn’t have $getField
-	// at all.) So we use a $switch instead.
-	mapAgg := bson.D{
-		{"$switch", bson.D{
-			{"branches", lo.Map(
-				slices.Collect(maps.Keys(mapping)),
-				func(key string, _ int) bson.D {
-					return bson.D{
-						{"case", bson.D{
-							{"$eq", bson.A{
-								key,
-								"$$numericKey",
-							}},
-						}},
-						{"then", mapping[key]},
-					}
-				},
-			)},
-		}},
-	}
-
 	return bson.D{
-		{"$arrayToObject", bson.D{
-			{"$map", bson.D{
-				{"input", bson.D{
-					{"$objectToArray", expr},
-				}},
-				{"in", bson.D{
-					{"$let", bson.D{
-						{"vars", bson.D{
-							{"numericKey", "$$this.k"},
-							{"value", "$$this.v"},
-						}},
-						{"in", bson.D{
-							{"k", mapAgg},
-							{"v", "$$value"},
-						}},
-					}},
-				}},
-			}},
-		}},
+		{"$arrayToObject", mslices.Of(lo.Map(
+			fieldNames,
+			func(fieldName string, _ int) [2]string {
+				return [...]string{fieldName, docExpr + "." + fieldName}
+			},
+		))},
 	}
 }

dockey/raw.go

@@ -23,6 +23,7 @@ func ExtractTrueDocKeyFromDoc(
 
 	var dk bson.D
 	for _, field := range fieldNames {
+		var val bson.RawValue
 
 		// This is how sharding routes documents: it always
 		// splits on the dot and looks deeply into the document.
@@ -31,13 +32,13 @@ func ExtractTrueDocKeyFromDoc(
 
 		if errors.Is(err, bsoncore.ErrElementNotFound) || errors.As(err, &bsoncore.InvalidDepthTraversalError{}) {
 			// If the document lacks a value for this field
-			// then don’t add it to the document key.
-			continue
-		} else if err == nil {
-			dk = append(dk, bson.E{field, val})
-		} else {
+			// then make it null in the document key.
+			val = bson.RawValue{Type: bson.TypeNull}
+		} else if err != nil {
 			return nil, errors.Wrapf(err, "extracting doc key field %#q from doc %+v", field, doc)
 		}
+
+		dk = append(dk, bson.E{field, val})
 	}
 
 	docKey, err := bson.Marshal(dk)

dockey/test/cases.go

@@ -52,6 +52,7 @@ var TestCases = []TestCase{
 		},
 		DocKey: bson.D{
 			{"_id", "ccc"},
+			{"foo.bar.baz", nil},
 		},
 	},
 	{
@@ -73,6 +74,7 @@ var TestCases = []TestCase{
 		},
 		DocKey: bson.D{
 			{"_id", "eee"},
+			{"foo.bar.baz", nil},
 		},
 	},
 }

internal/util/sharding.go

@@ -7,6 +7,8 @@ import (
 	"github.com/pkg/errors"
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo/writeconcern"
 )
 
 const (
@@ -52,3 +54,31 @@ func GetShardKey(
 
 	return option.Some(key), nil
 }
+
+// Used in tests:
+func DisableBalancing(ctx context.Context, coll *mongo.Collection) error {
+	client := coll.Database().Client()
+	configDB := client.Database("config")
+
+	ns := FullName(coll)
+
+	_, err := configDB.
+		Collection(
+			"collections",
+			options.Collection().SetWriteConcern(writeconcern.Majority()),
+		).
+		UpdateOne(
+			ctx,
+			bson.D{{"_id", ns}},
+			bson.D{
+				{"$set", bson.D{
+					{"noBalance", true},
+				}},
+			},
+		)
+
+	if err != nil {
+		return errors.Wrapf(err, "disabling %#q’s shard balancing", ns)
+	}
+	return nil
+}

internal/verifier/migration_verifier_test.go

@@ -162,70 +162,115 @@ func (suite *IntegrationTestSuite) TestVerifier_Dotted_Shard_Key() {
 	dbName := suite.DBNameForTest()
 	collName := "coll"
 
-	verifier := suite.BuildVerifier()
-	srcColl := verifier.srcClient.Database(dbName).Collection(collName)
-	dstColl := verifier.dstClient.Database(dbName).Collection(collName)
-
 	docs := []bson.D{
-		// NB: It’s not possible to test with /foo.bar.baz or /foo/bar.baz
-		// because the server (as of v8) will route those to the same shard.
-		// Only /foo/bar/baz actually gets routed as per the shard key, so
-		// only that value can create a duplicate ID.
-		{{"_id", 33}, {"foo", bson.D{{"bar", bson.D{{"baz", 100}}}}}},
-		{{"_id", 33}, {"foo", bson.D{{"bar", bson.D{{"baz", 200}}}}}},
+		{{"_id", 33}, {"foo", bson.D{{"bar", 100}}}},
+		{{"_id", 33}, {"foo", bson.D{{"bar", 200}}}},
 	}
 
 	shardKey := bson.D{
-		{"foo.bar.baz", 1},
+		{"foo.bar", 1},
 	}
 
-	for _, coll := range mslices.Of(srcColl, dstColl) {
-		db := coll.Database()
-		client := db.Client()
+	keyField := shardKey[0].Key
+	splitKey := bson.D{{keyField, 150}}
 
-		// For sharded, pre-v8 clusters we need to create the collection first.
-		require.NoError(db.CreateCollection(ctx, coll.Name()))
-		require.NoError(client.Database("admin").RunCommand(
-			ctx,
-			bson.D{{"enableSharding", db.Name()}},
-		).Err())
+	clientsWithLabels := []struct {
+		label  string
+		client *mongo.Client
+	}{
+		{"source", suite.srcMongoClient},
+		{"destination", suite.dstMongoClient},
+	}
 
-		require.NoError(client.Database("admin").RunCommand(
-			ctx,
-			bson.D{
-				{"shardCollection", db.Name() + "." + coll.Name()},
-				{"key", shardKey},
-			},
-		).Err(),
-		)
+	for _, clientWithLabel := range clientsWithLabels {
+		client := clientWithLabel.client
+		clientLabel := clientWithLabel.label
+
+		db := client.Database(dbName)
+		coll := db.Collection(collName)
+
+		// The DB shouldn’t exist anyway, but just in case.
+		require.NoError(db.Drop(ctx), "should drop database")
 
 		shardIds := getShardIds(suite.T(), client)
 
 		admin := client.Database("admin")
-		keyField := "foo.bar.baz"
-		splitKey := bson.D{{keyField, 150}}
+
+		require.NoError(admin.RunCommand(
+			ctx,
+			bson.D{
+				{"enableSharding", db.Name()},
+			},
+		).Err())
 
 		require.NoError(
-			admin.RunCommand(ctx, bson.D{
-				{"split", db.Name() + "." + coll.Name()},
-				{"middle", splitKey},
-			}).Err(),
+			admin.RunCommand(
+				ctx,
+				bson.D{
+					{"shardCollection", FullName(coll)},
+					{"key", shardKey},
+				},
+			).Err(),
+			"should shard collection on %s",
+			clientLabel,
 		)
 
 		require.NoError(
-			admin.RunCommand(ctx, bson.D{
-				{"moveChunk", db.Name() + "." + coll.Name()},
-				{"to", shardIds[0]},
-				{"find", bson.D{{keyField, 149}}},
-			}).Err(),
+			util.DisableBalancing(ctx, coll),
+			"should disable %#q’s balancing on %s",
+			FullName(coll),
+			clientLabel,
 		)
 
 		require.NoError(
 			admin.RunCommand(ctx, bson.D{
-				{"moveChunk", db.Name() + "." + coll.Name()},
-				{"to", shardIds[1]},
-				{"find", bson.D{{keyField, 151}}},
+				{"split", FullName(coll)},
+				{"middle", splitKey},
 			}).Err(),
+			"should split on %s",
+			clientLabel,
+		)
+
+		require.Eventually(
+			func() bool {
+				err := admin.RunCommand(ctx, bson.D{
+					{"moveChunk", FullName(coll)},
+					{"find", bson.D{{keyField, 149}}},
+					{"to", shardIds[0]},
+					{"_waitForDelete", true},
+				}).Err()
+
+				if err != nil {
+					suite.T().Logf("Failed to move %s’s lower chunk to shard %#q: %v", clientLabel, shardIds[0], err)
+					return false
+				}
+
+				return true
+			},
+			5*time.Minute,
+			time.Second,
+			"Should move lower chunk to the 1st shard",
+		)
+
+		require.Eventually(
+			func() bool {
+				err := admin.RunCommand(ctx, bson.D{
+					{"moveChunk", FullName(coll)},
+					{"find", bson.D{{keyField, 151}}},
+					{"to", shardIds[1]},
+					{"_waitForDelete", true},
+				}).Err()
+
+				if err != nil {
+					suite.T().Logf("Failed to move %s’s upper chunk to shard %#q: %v", clientLabel, shardIds[1], err)
+					return false
+				}
+
+				return true
+			},
+			5*time.Minute,
+			time.Second,
+			"Should move upper chunk to the 2nd shard",
 		)
 
 		_, err := coll.InsertMany(ctx, lo.ToAnySlice(lo.Shuffle(docs)))
@@ -252,6 +297,7 @@ func (suite *IntegrationTestSuite) TestVerifier_Dotted_Shard_Key() {
 		},
 	}
 
+	verifier := suite.BuildVerifier()
 	results, docCount, _, err := verifier.FetchAndCompareDocuments(ctx, 0, task)
 	require.NoError(err, "should fetch & compare")
 	assert.EqualValues(suite.T(), len(docs), docCount, "expected # of docs")