Fix type checks for update pipelines

GODRIVER-1231

Change-Id: Ie60d386efce586d0a271b487698fd0bfd13cd74c

Author: Divjot Arora

mongo/collection.go

@@ -430,8 +430,8 @@ func (coll *Collection) DeleteMany(ctx context.Context, filter interface{},
 	return coll.delete(ctx, filter, false, rrMany, opts...)
 }
 
-func (coll *Collection) updateOrReplace(ctx context.Context, filter bsoncore.Document, update interface{}, multi bool, expectedRr returnResult,
-	opts ...*options.UpdateOptions) (*UpdateResult, error) {
+func (coll *Collection) updateOrReplace(ctx context.Context, filter bsoncore.Document, update interface{}, multi bool,
+	expectedRr returnResult, checkDollarKey bool, opts ...*options.UpdateOptions) (*UpdateResult, error) {
 
 	if ctx == nil {
 		ctx = context.Background()
@@ -441,16 +441,11 @@ func (coll *Collection) updateOrReplace(ctx context.Context, filter bsoncore.Doc
 	uidx, updateDoc := bsoncore.AppendDocumentStart(nil)
 	updateDoc = bsoncore.AppendDocumentElement(updateDoc, "q", filter)
 
-	switch update.(type) {
-	case bsoncore.Document:
-		updateDoc = bsoncore.AppendDocumentElement(updateDoc, "u", update.(bsoncore.Document))
-	default:
-		u, err := transformUpdateValue(coll.registry, update, true)
-		if err != nil {
-			return nil, err
-		}
-		updateDoc = bsoncore.AppendValueElement(updateDoc, "u", u)
+	u, err := transformUpdateValue(coll.registry, update, checkDollarKey)
+	if err != nil {
+		return nil, err
 	}
+	updateDoc = bsoncore.AppendValueElement(updateDoc, "u", u)
 	if multi {
 		updateDoc = bsoncore.AppendBooleanElement(updateDoc, "multi", multi)
 	}
@@ -482,7 +477,7 @@ func (coll *Collection) updateOrReplace(ctx context.Context, filter bsoncore.Doc
 		defer sess.EndSession()
 	}
 
-	err := coll.client.validSession(sess)
+	err = coll.client.validSession(sess)
 	if err != nil {
 		return nil, err
 	}
@@ -546,7 +541,7 @@ func (coll *Collection) UpdateOne(ctx context.Context, filter interface{}, updat
 		return nil, err
 	}
 
-	return coll.updateOrReplace(ctx, f, update, false, rrOne, opts...)
+	return coll.updateOrReplace(ctx, f, update, false, rrOne, true, opts...)
 }
 
 // UpdateMany updates multiple documents in the collection.
@@ -562,7 +557,7 @@ func (coll *Collection) UpdateMany(ctx context.Context, filter interface{}, upda
 		return nil, err
 	}
 
-	return coll.updateOrReplace(ctx, f, update, true, rrMany, opts...)
+	return coll.updateOrReplace(ctx, f, update, true, rrMany, true, opts...)
 }
 
 // ReplaceOne replaces a single document in the collection.
@@ -596,7 +591,7 @@ func (coll *Collection) ReplaceOne(ctx context.Context, filter interface{},
 		updateOptions = append(updateOptions, uOpts)
 	}
 
-	return coll.updateOrReplace(ctx, f, r, false, rrOne, updateOptions...)
+	return coll.updateOrReplace(ctx, f, r, false, rrOne, false, updateOptions...)
 }
 
 // Aggregate runs an aggregation framework pipeline.

mongo/collection_internal_test.go

@@ -12,10 +12,6 @@ import (
 	"fmt"
 	"os"
 	"testing"
-
-	"go.mongodb.org/mongo-driver/mongo/options"
-	"go.mongodb.org/mongo-driver/x/bsonx"
-
 	"time"
 
 	"github.com/google/go-cmp/cmp"
@@ -25,9 +21,12 @@ import (
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/event"
 	"go.mongodb.org/mongo-driver/internal/testutil"
+	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.mongodb.org/mongo-driver/mongo/readconcern"
 	"go.mongodb.org/mongo-driver/mongo/readpref"
 	"go.mongodb.org/mongo-driver/mongo/writeconcern"
+	"go.mongodb.org/mongo-driver/x/bsonx"
+	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
 var impossibleWriteConcern = writeconcern.New(writeconcern.W(50), writeconcern.WTimeout(time.Second))
@@ -2221,3 +2220,44 @@ func TestCollection_BulkWrite(t *testing.T) {
 		}
 	})
 }
+
+// test special types that should be converted to a document for updates even though the underlying type is a
+// slice/array
+func TestCollection_Update_SpecialSliceTypes(t *testing.T) {
+	doc := bson.D{{"$set", bson.D{{"x", 2}}}}
+	docBytes, err := bson.Marshal(doc)
+	require.NoError(t, err, "error getting document bytes: %v", err)
+	xUpdate := bsonx.Doc{{"x", bsonx.Int32(2)}}
+	xDoc := bsonx.Doc{{"$set", bsonx.Document(xUpdate)}}
+
+	testCases := []struct {
+		name   string
+		update interface{}
+	}{
+		{"bsoncore Document", bsoncore.Document(docBytes)},
+		{"bson Raw", bson.Raw(docBytes)},
+		{"bson D", doc},
+		{"bsonx Document", xDoc},
+		{"byte slice", docBytes},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			coll := createTestCollection(t, nil, &tc.name)
+			defer func() {
+				_ = coll.Drop(ctx)
+			}()
+
+			insertedDoc := bson.D{{"x", 1}}
+			_, err = coll.InsertOne(ctx, insertedDoc)
+			require.NoError(t, err, "error inserting document: %v", err)
+
+			res, err := coll.UpdateOne(ctx, insertedDoc, tc.update)
+			require.NoError(t, err, "error updating document: %v", err)
+			require.Equal(t, int64(1), res.MatchedCount,
+				"matched count mismatch; expected %d, got %d", 1, res.MatchedCount)
+			require.Equal(t, int64(1), res.ModifiedCount,
+				"modified count mismatch; expected %d, got %d", 1, res.ModifiedCount)
+		})
+	}
+}

mongo/mongo.go

@@ -339,13 +339,34 @@ func transformUpdateValue(registry *bsoncodec.Registry, update interface{}, chec
 	switch t := update.(type) {
 	case nil:
 		return u, ErrNilDocument
-	case primitive.D:
+	case primitive.D, bsonx.Doc:
 		u.Type = bsontype.EmbeddedDocument
 		u.Data, err = transformBsoncoreDocument(registry, update)
 		if err != nil {
 			return u, err
 		}
 
+		if checkDocDollarKey {
+			err = ensureDollarKeyv2(u.Data)
+		}
+		return u, err
+	case bson.Raw:
+		u.Type = bsontype.EmbeddedDocument
+		u.Data = t
+		if checkDocDollarKey {
+			err = ensureDollarKeyv2(u.Data)
+		}
+		return u, err
+	case bsoncore.Document:
+		u.Type = bsontype.EmbeddedDocument
+		u.Data = t
+		if checkDocDollarKey {
+			err = ensureDollarKeyv2(u.Data)
+		}
+		return u, err
+	case []byte:
+		u.Type = bsontype.EmbeddedDocument
+		u.Data = t
 		if checkDocDollarKey {
 			err = ensureDollarKeyv2(u.Data)
 		}