GODRIVER-1363 fix indexes for writeErrors with batches and filter ret… (#362)

Author: iwysiu

mongo/collection.go

@@ -293,7 +293,25 @@ func (coll *Collection) insert(ctx context.Context, documents []interface{},
 	}
 	op = op.Retry(retry)
 
-	return result, op.Execute(ctx)
+	err = op.Execute(ctx)
+	wce, ok := err.(driver.WriteCommandError)
+	if !ok {
+		return result, err
+	}
+
+	// remove the ids that had writeErrors from result
+	for i, we := range wce.WriteErrors {
+		// i indexes have been removed before the current error, so the index is we.Index-i
+		idIndex := int(we.Index) - i
+		// if the insert is ordered, nothing after the error was inserted
+		if imo.Ordered == nil || *imo.Ordered {
+			result = result[:idIndex]
+			break
+		}
+		result = append(result[:idIndex], result[idIndex+1:]...)
+	}
+
+	return result, err
 }
 
 // InsertOne executes an insert command to insert a single document into the collection.
@@ -367,6 +385,7 @@ func (coll *Collection) InsertMany(ctx context.Context, documents []interface{},
 			nil,
 		})
 	}
+
 	return imResult, BulkWriteException{
 		WriteErrors:       bwErrors,
 		WriteConcernError: writeException.WriteConcernError,

mongo/integration/collection_test.go

@@ -164,6 +164,74 @@ func TestCollection(t *testing.T) {
 				})
 			}
 		})
+		mt.Run("return only inserted ids", func(mt *mtest.T) {
+			id := int32(11)
+			docs := []interface{}{
+				bson.D{{"_id", id}},
+				bson.D{{"_id", id}},
+				bson.D{{"x", 6}},
+				bson.D{{"_id", id}},
+			}
+
+			testCases := []struct {
+				name        string
+				ordered     bool
+				numInserted int
+				numErrors   int
+			}{
+				{"unordered", false, 2, 2},
+				{"ordered", true, 1, 1},
+			}
+			for _, tc := range testCases {
+				mt.Run(tc.name, func(mt *mtest.T) {
+					res, err := mt.Coll.InsertMany(mtest.Background, docs, options.InsertMany().SetOrdered(tc.ordered))
+
+					assert.Equal(mt, tc.numInserted, len(res.InsertedIDs), "expected %v inserted IDs, got %v", tc.numInserted, len(res.InsertedIDs))
+					assert.Equal(mt, id, res.InsertedIDs[0], "expected inserted ID %v, got %v", id, res.InsertedIDs[0])
+					if tc.numInserted > 1 {
+						assert.NotNil(mt, res.InsertedIDs[1], "expected ID but got nil")
+					}
+
+					we, ok := err.(mongo.BulkWriteException)
+					assert.True(mt, ok, "expected error type %T, got %T", mongo.BulkWriteException{}, err)
+					numErrors := len(we.WriteErrors)
+					assert.Equal(mt, tc.numErrors, numErrors, "expected %v write errors, got %v", tc.numErrors, numErrors)
+					gotCode := we.WriteErrors[0].Code
+					assert.Equal(mt, errorDuplicateKey, gotCode, "expected error code %v, got %v", errorDuplicateKey, gotCode)
+				})
+			}
+		})
+		mt.Run("writeError index", func(mt *mtest.T) {
+			// TODO(GODRIVER-425): remove this as part a larger project to
+			// refactor integration and other longrunning tasks.
+			if os.Getenv("EVR_TASK_ID") == "" {
+				mt.Skip("skipping long running integration test outside of evergreen")
+			}
+
+			// force multiple batches
+			numDocs := 700000
+			var docs []interface{}
+			for i := 0; i < numDocs; i++ {
+				d := bson.D{
+					{"a", int32(i)},
+					{"b", int32(i * 2)},
+					{"c", int32(i * 3)},
+				}
+				docs = append(docs, d)
+			}
+			repeated := bson.D{{"_id", int32(11)}}
+			docs = append(docs, repeated, repeated)
+
+			_, err := mt.Coll.InsertMany(context.Background(), docs)
+			assert.NotNil(mt, err, "expected InsertMany error, got nil")
+
+			we, ok := err.(mongo.BulkWriteException)
+			assert.True(mt, ok, "expected error type %T, got %T", mongo.BulkWriteException{}, err)
+			numErrors := len(we.WriteErrors)
+			assert.Equal(mt, 1, numErrors, "expected 1 write error, got %v", numErrors)
+			gotIndex := we.WriteErrors[0].Index
+			assert.Equal(mt, numDocs+1, gotIndex, "expected index %v, got %v", numDocs+1, gotIndex)
+		})
 		wcCollOpts := options.Collection().SetWriteConcern(impossibleWc)
 		wcTestOpts := mtest.NewOptions().CollectionOptions(wcCollOpts).Topologies(mtest.ReplicaSet)
 		mt.RunOpts("write concern error", wcTestOpts, func(mt *mtest.T) {

x/mongo/driver/operation.go

@@ -305,6 +305,7 @@ func (op Operation) Execute(ctx context.Context, scratch []byte) error {
 	}
 	batching := op.Batches.Valid()
 	retryEnabled := op.RetryMode != nil && op.RetryMode.Enabled()
+	currIndex := 0
 	for {
 		if batching {
 			targetBatchSize := desc.MaxDocumentSize
@@ -430,6 +431,13 @@ func (op Operation) Execute(ctx context.Context, scratch []byte) error {
 				}
 				continue
 			}
+
+			if batching && len(tt.WriteErrors) > 0 && currIndex > 0 {
+				for i := range tt.WriteErrors {
+					tt.WriteErrors[i].Index += int64(currIndex)
+				}
+			}
+
 			// If batching is enabled and either ordered is the default (which is true) or
 			// explicitly set to true and we have write errors, return the errors.
 			if batching && (op.Batches.Ordered == nil || *op.Batches.Ordered == true) && len(tt.WriteErrors) > 0 {
@@ -526,6 +534,7 @@ func (op Operation) Execute(ctx context.Context, scratch []byte) error {
 					retries = 1
 				}
 			}
+			currIndex += len(op.Batches.Current)
 			op.Batches.ClearBatch()
 			continue
 		}