GODRIVER-2232 Allow bsoncore.Array as type for aggregation pipelines (#807)

Author: benjirewis

mongo/mongo.go

@@ -210,20 +210,43 @@ func transformAggregatePipeline(registry *bsoncodec.Registry, pipeline interface
 			return nil, false, fmt.Errorf("can only transform slices and arrays into aggregation pipelines, but got %v", val.Kind())
 		}
 
-		aidx, arr := bsoncore.AppendArrayStart(nil)
 		var hasOutputStage bool
 		valLen := val.Len()
 
+		switch t := pipeline.(type) {
 		// Explicitly forbid non-empty pipelines that are semantically single documents
 		// and are implemented as slices.
-		switch t := pipeline.(type) {
 		case bson.D, bson.Raw, bsoncore.Document:
 			if valLen > 0 {
 				return nil, false,
 					fmt.Errorf("%T is not an allowed pipeline type as it represents a single document. Use bson.A or mongo.Pipeline instead", t)
 			}
+		// bsoncore.Arrays do not need to be transformed. Only check validity and presence of output stage.
+		case bsoncore.Array:
+			if err := t.Validate(); err != nil {
+				return nil, false, err
+			}
+
+			values, err := t.Values()
+			if err != nil {
+				return nil, false, err
+			}
+
+			numVals := len(values)
+			if numVals == 0 {
+				return bsoncore.Document(t), false, nil
+			}
+
+			// If not empty, check if first value of the last stage is $out or $merge.
+			if lastStage, ok := values[numVals-1].DocumentOK(); ok {
+				if elem, err := lastStage.IndexErr(0); err == nil && (elem.Key() == "$out" || elem.Key() == "$merge") {
+					hasOutputStage = true
+				}
+			}
+			return bsoncore.Document(t), hasOutputStage, nil
 		}
 
+		aidx, arr := bsoncore.AppendArrayStart(nil)
 		for idx := 0; idx < valLen; idx++ {
 			doc, err := transformBsoncoreDocument(registry, val.Index(idx).Interface(), true, fmt.Sprintf("pipeline stage :%v", idx))
 			if err != nil {

mongo/mongo_test.go

@@ -120,16 +120,43 @@ func TestMongoHelpers(t *testing.T) {
 		})
 	})
 	t.Run("transform aggregate pipeline", func(t *testing.T) {
+		// []byte of [{{"$limit", 12345}}]
 		index, arr := bsoncore.AppendArrayStart(nil)
 		dindex, arr := bsoncore.AppendDocumentElementStart(arr, "0")
 		arr = bsoncore.AppendInt32Element(arr, "$limit", 12345)
 		arr, _ = bsoncore.AppendDocumentEnd(arr, dindex)
 		arr, _ = bsoncore.AppendArrayEnd(arr, index)
 
+		// []byte of {{"x", 1}}
 		index, doc := bsoncore.AppendDocumentStart(nil)
 		doc = bsoncore.AppendInt32Element(doc, "x", 1)
 		doc, _ = bsoncore.AppendDocumentEnd(doc, index)
 
+		// bsoncore.Array of [{{"$merge", {}}}]
+		mergeStage := bsoncore.NewDocumentBuilder().
+			StartDocument("$merge").
+			FinishDocument().
+			Build()
+		arrMergeStage := bsoncore.NewArrayBuilder().AppendDocument(mergeStage).Build()
+
+		fooStage := bsoncore.NewDocumentBuilder().AppendString("foo", "bar").Build()
+		bazStage := bsoncore.NewDocumentBuilder().AppendString("baz", "qux").Build()
+		outStage := bsoncore.NewDocumentBuilder().AppendString("$out", "myColl").Build()
+
+		// bsoncore.Array of [{{"foo", "bar"}}, {{"baz", "qux"}}, {{"$out", "myColl"}}]
+		arrOutStage := bsoncore.NewArrayBuilder().
+			AppendDocument(fooStage).
+			AppendDocument(bazStage).
+			AppendDocument(outStage).
+			Build()
+
+		// bsoncore.Array of [{{"foo", "bar"}}, {{"$out", "myColl"}}, {{"baz", "qux"}}]
+		arrMiddleOutStage := bsoncore.NewArrayBuilder().
+			AppendDocument(fooStage).
+			AppendDocument(outStage).
+			AppendDocument(bazStage).
+			Build()
+
 		testCases := []struct {
 			name           string
 			pipeline       interface{}
@@ -388,6 +415,46 @@ func TestMongoHelpers(t *testing.T) {
 				false,
 				nil,
 			},
+			{
+				"bsoncore.Array/success",
+				bsoncore.Array(arr),
+				bson.A{
+					bson.D{{"$limit", int32(12345)}},
+				},
+				false,
+				nil,
+			},
+			{
+				"bsoncore.Array/mergeStage",
+				arrMergeStage,
+				bson.A{
+					bson.D{{"$merge", bson.D{}}},
+				},
+				true,
+				nil,
+			},
+			{
+				"bsoncore.Array/outStage",
+				arrOutStage,
+				bson.A{
+					bson.D{{"foo", "bar"}},
+					bson.D{{"baz", "qux"}},
+					bson.D{{"$out", "myColl"}},
+				},
+				true,
+				nil,
+			},
+			{
+				"bsoncore.Array/middleOutStage",
+				arrMiddleOutStage,
+				bson.A{
+					bson.D{{"foo", "bar"}},
+					bson.D{{"$out", "myColl"}},
+					bson.D{{"baz", "qux"}},
+				},
+				false,
+				nil,
+			},
 		}
 
 		for _, tc := range testCases {