GODRIVER-1782 correctly mark aggregation pipeline with $out (#537)

Author: iwysiu

mongo/integration/crud_prose_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 
 	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/bsontype"
 	"go.mongodb.org/mongo-driver/internal/testutil/assert"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/integration/mtest"
@@ -145,6 +146,14 @@ func TestHintErrors(t *testing.T) {
 	})
 }
 
+type testValueMarshaler struct {
+	val []bson.D
+}
+
+func (tvm testValueMarshaler) MarshalBSONValue() (bsontype.Type, []byte, error) {
+	return bson.MarshalValue(tvm.val)
+}
+
 func TestAggregatePrimaryPreferredReadPreference(t *testing.T) {
 	primaryPrefClientOpts := options.Client().
 		SetWriteConcern(mtest.MajorityWc).
@@ -155,39 +164,62 @@ func TestAggregatePrimaryPreferredReadPreference(t *testing.T) {
 		MinServerVersion("4.1.0") // Consistent with tests in aggregate-out-readConcern.json
 
 	mt := mtest.New(t, mtOpts)
-	mt.Run("aggregate $out with non-primary read ppreference", func(mt *mtest.T) {
+	mt.Run("aggregate $out with non-primary read preference", func(mt *mtest.T) {
 		doc, err := bson.Marshal(bson.D{
 			{"_id", 1},
 			{"x", 11},
 		})
 		assert.Nil(mt, err, "Marshal error: %v", err)
-		_, err = mt.Coll.InsertOne(mtest.Background, doc)
-		assert.Nil(mt, err, "InsertOne error: %v", err)
-
-		mt.ClearEvents()
 		outputCollName := "aggregate-read-pref-primary-preferred-output"
-		outStage := bson.D{
-			{"$out", outputCollName},
+		testCases := []struct {
+			name     string
+			pipeline interface{}
+		}{
+			{
+				"pipeline",
+				mongo.Pipeline{bson.D{{"$out", outputCollName}}},
+			},
+			{
+				"doc slice",
+				[]bson.D{{{"$out", outputCollName}}},
+			},
+			{
+				"bson a",
+				bson.A{bson.D{{"$out", outputCollName}}},
+			},
+			{
+				"valueMarshaler",
+				testValueMarshaler{[]bson.D{{{"$out", outputCollName}}}},
+			},
+		}
+		for _, tc := range testCases {
+			mt.Run(tc.name, func(mt *mtest.T) {
+				_, err = mt.Coll.InsertOne(mtest.Background, doc)
+				assert.Nil(mt, err, "InsertOne error: %v", err)
+
+				mt.ClearEvents()
+
+				cursor, err := mt.Coll.Aggregate(mtest.Background, tc.pipeline)
+				assert.Nil(mt, err, "Aggregate error: %v", err)
+				_ = cursor.Close(mtest.Background)
+
+				// Assert that the output collection contains the document we expect.
+				outputColl := mt.CreateCollection(mtest.Collection{Name: outputCollName}, false)
+				cursor, err = outputColl.Find(mtest.Background, bson.D{})
+				assert.Nil(mt, err, "Find error: %v", err)
+				defer cursor.Close(mtest.Background)
+
+				assert.True(mt, cursor.Next(mtest.Background), "expected Next to return true, got false")
+				assert.True(mt, bytes.Equal(doc, cursor.Current), "expected document %s, got %s", bson.Raw(doc), cursor.Current)
+				assert.False(mt, cursor.Next(mtest.Background), "unexpected document returned by Find: %s", cursor.Current)
+
+				// Assert that no read preference was sent to the server.
+				evt := mt.GetStartedEvent()
+				assert.Equal(mt, "aggregate", evt.CommandName, "expected command 'aggregate', got '%s'", evt.CommandName)
+				_, err = evt.Command.LookupErr("$readPreference")
+				assert.NotNil(mt, err, "expected command %s to not contain $readPreference", evt.Command)
+			})
 		}
-		cursor, err := mt.Coll.Aggregate(mtest.Background, mongo.Pipeline{outStage})
-		assert.Nil(mt, err, "Aggregate error: %v", err)
-		_ = cursor.Close(mtest.Background)
-
-		// Assert that the output collection contains the document we expect.
-		outputColl := mt.CreateCollection(mtest.Collection{Name: outputCollName}, false)
-		cursor, err = outputColl.Find(mtest.Background, bson.D{})
-		assert.Nil(mt, err, "Find error: %v", err)
-		defer cursor.Close(mtest.Background)
-
-		assert.True(mt, cursor.Next(mtest.Background), "expected Next to return true, got false")
-		assert.True(mt, bytes.Equal(doc, cursor.Current), "expected document %s, got %s", bson.Raw(doc), cursor.Current)
-		assert.False(mt, cursor.Next(mtest.Background), "unexpected document returned by Find: %s", cursor.Current)
-
-		// Assert that no read preference was sent to the server.
-		evt := mt.GetStartedEvent()
-		assert.Equal(mt, "aggregate", evt.CommandName, "expected command 'aggregate', got '%s'", evt.CommandName)
-		_, err = evt.Command.LookupErr("$readPreference")
-		assert.NotNil(mt, err, "expected command %s to not contain $readPreference", evt.Command)
 	})
 }
 

mongo/mongo.go

@@ -306,11 +306,13 @@ func transformAggregatePipelinev2(registry *bsoncodec.Registry, pipeline interfa
 
 		var hasOutputStage bool
 		pipelineDoc := bsoncore.Document(val)
-		if _, err := pipelineDoc.LookupErr("$out"); err == nil {
-			hasOutputStage = true
-		}
-		if _, err := pipelineDoc.LookupErr("$merge"); err == nil {
-			hasOutputStage = true
+		values, _ := pipelineDoc.Values()
+		if pipelineLen := len(values); pipelineLen > 0 {
+			if finalDoc, ok := values[pipelineLen-1].DocumentOK(); ok {
+				if elem, err := finalDoc.IndexErr(0); err == nil && (elem.Key() == "$out" || elem.Key() == "$merge") {
+					hasOutputStage = true
+				}
+			}
 		}
 
 		return pipelineDoc, hasOutputStage, nil