REP-6465: Make dockey always include shard key fields even if empty.

This amends PR #127 so that, if a document lacks a shard key field, the
generated document key will include that field with a null value. This
matches the server’s internal logic and seems safer besides.

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},
 		},
 	},
 }