GODRIVER-1696 CRUD update/replace bug fixes

Divjot Arora · Divjot Arora · commit 1ce074c35c7a · 2020-07-17T18:13:14.000-04:00
diff --git a/mongo/bulk_write.go b/mongo/bulk_write.go
@@ -281,6 +281,7 @@ func createDeleteDoc(filter interface{}, collation *options.Collation, hint inte
 func (bw *bulkWrite) runUpdate(ctx context.Context, batch bulkWriteBatch) (operation.UpdateResult, error) {
 	docs := make([]bsoncore.Document, len(batch.models))
 	var hasHint bool
+	var hasArrayFilters bool
 	for i, model := range batch.models {
 		var doc bsoncore.Document
 		var err error
@@ -292,12 +293,14 @@ func (bw *bulkWrite) runUpdate(ctx context.Context, batch bulkWriteBatch) (opera
 			hasHint = hasHint || (converted.Hint != nil)
 		case *UpdateOneModel:
 			doc, err = createUpdateDoc(converted.Filter, converted.Update, converted.Hint, converted.ArrayFilters, converted.Collation, converted.Upsert, false,
-				false, bw.collection.registry)
+				true, bw.collection.registry)
 			hasHint = hasHint || (converted.Hint != nil)
+			hasArrayFilters = hasArrayFilters || (converted.ArrayFilters != nil)
 		case *UpdateManyModel:
 			doc, err = createUpdateDoc(converted.Filter, converted.Update, converted.Hint, converted.ArrayFilters, converted.Collation, converted.Upsert, true,
-				false, bw.collection.registry)
+				true, bw.collection.registry)
 			hasHint = hasHint || (converted.Hint != nil)
+			hasArrayFilters = hasArrayFilters || (converted.ArrayFilters != nil)
 		}
 		if err != nil {
 			return operation.UpdateResult{}, err
@@ -310,7 +313,8 @@ func (bw *bulkWrite) runUpdate(ctx context.Context, batch bulkWriteBatch) (opera
 		Session(bw.session).WriteConcern(bw.writeConcern).CommandMonitor(bw.collection.client.monitor).
 		ServerSelector(bw.selector).ClusterClock(bw.collection.client.clock).
 		Database(bw.collection.db.name).Collection(bw.collection.name).
-		Deployment(bw.collection.client.deployment).Crypt(bw.collection.client.crypt).Hint(hasHint)
+		Deployment(bw.collection.client.deployment).Crypt(bw.collection.client.crypt).Hint(hasHint).
+		ArrayFilters(hasArrayFilters)
 	if bw.ordered != nil {
 		op = op.Ordered(*bw.ordered)
 	}
@@ -350,6 +354,7 @@ func createUpdateDoc(
 	if err != nil {
 		return nil, err
 	}
+
 	updateDoc = bsoncore.AppendValueElement(updateDoc, "u", u)
 
 	if multi {
diff --git a/mongo/collection.go b/mongo/collection.go
@@ -550,7 +550,8 @@ func (coll *Collection) updateOrReplace(ctx context.Context, filter bsoncore.Doc
 		Session(sess).WriteConcern(wc).CommandMonitor(coll.client.monitor).
 		ServerSelector(selector).ClusterClock(coll.client.clock).
 		Database(coll.db.name).Collection(coll.name).
-		Deployment(coll.client.deployment).Crypt(coll.client.crypt).Hint(uo.Hint != nil)
+		Deployment(coll.client.deployment).Crypt(coll.client.crypt).Hint(uo.Hint != nil).
+		ArrayFilters(uo.ArrayFilters != nil)
 
 	if uo.BypassDocumentValidation != nil && *uo.BypassDocumentValidation {
 		op = op.BypassDocumentValidation(*uo.BypassDocumentValidation)
@@ -669,8 +670,8 @@ func (coll *Collection) ReplaceOne(ctx context.Context, filter interface{},
 		return nil, err
 	}
 
-	if elem, err := r.IndexErr(0); err == nil && strings.HasPrefix(elem.Key(), "$") {
-		return nil, errors.New("replacement document cannot contains keys beginning with '$")
+	if err := ensureNoDollarKey(r); err != nil {
+		return nil, err
 	}
 
 	updateOptions := make([]*options.UpdateOptions, 0, len(opts))
diff --git a/mongo/mongo.go b/mongo/mongo.go
@@ -253,6 +253,14 @@ func ensureDollarKeyv2(doc bsoncore.Document) error {
 	return nil
 }
 
+func ensureNoDollarKey(doc bsoncore.Document) error {
+	if elem, err := doc.IndexErr(0); err == nil && strings.HasPrefix(elem.Key(), "$") {
+		return errors.New("replacement document cannot contains keys beginning with '$")
+	}
+
+	return nil
+}
+
 func transformAggregatePipeline(registry *bsoncodec.Registry, pipeline interface{}) (bsonx.Arr, error) {
 	pipelineArr := bsonx.Arr{}
 	switch t := pipeline.(type) {
@@ -333,7 +341,12 @@ func transformAggregatePipelinev2(registry *bsoncodec.Registry, pipeline interfa
 	}
 }
 
-func transformUpdateValue(registry *bsoncodec.Registry, update interface{}, checkDocDollarKey bool) (bsoncore.Value, error) {
+func transformUpdateValue(registry *bsoncodec.Registry, update interface{}, dollarKeysAllowed bool) (bsoncore.Value, error) {
+	documentCheckerFunc := ensureDollarKeyv2
+	if !dollarKeysAllowed {
+		documentCheckerFunc = ensureNoDollarKey
+	}
+
 	var u bsoncore.Value
 	var err error
 	switch t := update.(type) {
@@ -346,42 +359,27 @@ func transformUpdateValue(registry *bsoncodec.Registry, update interface{}, chec
 			return u, err
 		}
 
-		if checkDocDollarKey {
-			err = ensureDollarKeyv2(u.Data)
-		}
-		return u, err
+		return u, documentCheckerFunc(u.Data)
 	case bson.Raw:
 		u.Type = bsontype.EmbeddedDocument
 		u.Data = t
-		if checkDocDollarKey {
-			err = ensureDollarKeyv2(u.Data)
-		}
-		return u, err
+		return u, documentCheckerFunc(u.Data)
 	case bsoncore.Document:
 		u.Type = bsontype.EmbeddedDocument
 		u.Data = t
-		if checkDocDollarKey {
-			err = ensureDollarKeyv2(u.Data)
-		}
-		return u, err
+		return u, documentCheckerFunc(u.Data)
 	case []byte:
 		u.Type = bsontype.EmbeddedDocument
 		u.Data = t
-		if checkDocDollarKey {
-			err = ensureDollarKeyv2(u.Data)
-		}
-		return u, err
+		return u, documentCheckerFunc(u.Data)
 	case bsoncodec.Marshaler:
 		u.Type = bsontype.EmbeddedDocument
 		u.Data, err = t.MarshalBSON()
 		if err != nil {
 			return u, err
 		}
 
-		if checkDocDollarKey {
-			err = ensureDollarKeyv2(u.Data)
-		}
-		return u, err
+		return u, documentCheckerFunc(u.Data)
 	case bsoncodec.ValueMarshaler:
 		u.Type, u.Data, err = t.MarshalBSONValue()
 		if err != nil {
@@ -403,10 +401,7 @@ func transformUpdateValue(registry *bsoncodec.Registry, update interface{}, chec
 				return u, err
 			}
 
-			if checkDocDollarKey {
-				err = ensureDollarKeyv2(u.Data)
-			}
-			return u, err
+			return u, documentCheckerFunc(u.Data)
 		}
 
 		u.Type = bsontype.Array
@@ -418,7 +413,7 @@ func transformUpdateValue(registry *bsoncodec.Registry, update interface{}, chec
 				return u, err
 			}
 
-			if err := ensureDollarKeyv2(doc); err != nil {
+			if err := documentCheckerFunc(doc); err != nil {
 				return u, err
 			}
 
diff --git a/x/mongo/driver/operation/update.go b/x/mongo/driver/operation/update.go
@@ -34,6 +34,7 @@ type Update struct {
 	database                 string
 	deployment               driver.Deployment
 	hint                     *bool
+	arrayFilters             *bool
 	selector                 description.ServerSelector
 	writeConcern             *writeconcern.WriteConcern
 	retry                    *driver.RetryMode
@@ -177,6 +178,11 @@ func (u *Update) command(dst []byte, desc description.SelectedServer) ([]byte, e
 			return nil, errUnacknowledgedHint
 		}
 	}
+	if u.arrayFilters != nil && *u.arrayFilters {
+		if desc.WireVersion == nil || !desc.WireVersion.Includes(6) {
+			return nil, errors.New("the 'arrayFilters' command parameter requires a minimum server wire version of 6")
+		}
+	}
 
 	return dst, nil
 }
@@ -204,6 +210,17 @@ func (u *Update) Hint(hint bool) *Update {
 	return u
 }
 
+// ArrayFilters is a flag to indicate that the update document contains an arrayFilters field. This option is only
+// supported on server versions 3.6 and higher. For servers < 3.6, the driver will return an error.
+func (u *Update) ArrayFilters(arrayFilters bool) *Update {
+	if u == nil {
+		u = new(Update)
+	}
+
+	u.arrayFilters = &arrayFilters
+	return u
+}
+
 // Ordered sets ordered. If true, when a write fails, the operation will return the error, when
 // false write failures do not stop execution of the operation.
 func (u *Update) Ordered(ordered bool) *Update {
