GODRIVER-1028 return correct error indexes for bulkWrite (mongodb#483)

Author: iwysiu

mongo/bulk_write.go

@@ -22,6 +22,7 @@ import (
 type bulkWriteBatch struct {
 	models   []WriteModel
 	canRetry bool
+	indexes  []int
 }
 
 // bulkWrite perfoms a bulkwrite operation
@@ -52,7 +53,6 @@ func (bw *bulkWrite) execute(ctx context.Context) error {
 	}
 
 	var lastErr error
-	var opIndex int64 // the operation index for the upsertedIDs map
 	continueOnError := !ordered
 	for _, batch := range batches {
 		if len(batch.models) == 0 {
@@ -66,13 +66,10 @@ func (bw *bulkWrite) execute(ctx context.Context) error {
 
 		batchRes, batchErr, err := bw.runBatch(ctx, batch)
 
-		bw.mergeResults(batchRes, opIndex)
+		bw.mergeResults(batchRes)
 
 		bwErr.WriteConcernError = batchErr.WriteConcernError
 		bwErr.Labels = append(bwErr.Labels, batchErr.Labels...)
-		for i := range batchErr.WriteErrors {
-			batchErr.WriteErrors[i].Index = batchErr.WriteErrors[i].Index + int(opIndex)
-		}
 
 		bwErr.WriteErrors = append(bwErr.WriteErrors, batchErr.WriteErrors...)
 
@@ -89,8 +86,6 @@ func (bw *bulkWrite) execute(ctx context.Context) error {
 		if err != nil {
 			lastErr = err
 		}
-
-		opIndex += int64(len(batch.models))
 	}
 
 	bw.result.MatchedCount -= bw.result.UpsertedCount
@@ -151,16 +146,18 @@ func (bw *bulkWrite) runBatch(ctx context.Context, batch bulkWriteBatch) (BulkWr
 		batchRes.ModifiedCount = int64(res.NModified)
 		batchRes.UpsertedCount = int64(len(res.Upserted))
 		for _, upsert := range res.Upserted {
-			batchRes.UpsertedIDs[upsert.Index] = upsert.ID
+			batchRes.UpsertedIDs[int64(batch.indexes[upsert.Index])] = upsert.ID
 		}
 	}
 
 	batchErr.WriteErrors = make([]BulkWriteError, 0, len(writeErrors))
 	convWriteErrors := writeErrorsFromDriverWriteErrors(writeErrors)
 	for _, we := range convWriteErrors {
+		request := batch.models[we.Index]
+		we.Index = batch.indexes[we.Index]
 		batchErr.WriteErrors = append(batchErr.WriteErrors, BulkWriteError{
 			WriteError: we,
-			Request:    batch.models[we.Index],
+			Request:    request,
 		})
 	}
 	return batchRes, batchErr, nil
@@ -401,18 +398,23 @@ func createBatches(models []WriteModel, ordered bool) []bulkWriteBatch {
 	batches[updateOneCommand].canRetry = true
 
 	// TODO(GODRIVER-1157): fix batching once operation retryability is fixed
-	for _, model := range models {
+	for i, model := range models {
 		switch model.(type) {
 		case *InsertOneModel:
 			batches[insertCommand].models = append(batches[insertCommand].models, model)
+			batches[insertCommand].indexes = append(batches[insertCommand].indexes, i)
 		case *DeleteOneModel:
 			batches[deleteOneCommand].models = append(batches[deleteOneCommand].models, model)
+			batches[deleteOneCommand].indexes = append(batches[deleteOneCommand].indexes, i)
 		case *DeleteManyModel:
 			batches[deleteManyCommand].models = append(batches[deleteManyCommand].models, model)
+			batches[deleteManyCommand].indexes = append(batches[deleteManyCommand].indexes, i)
 		case *ReplaceOneModel, *UpdateOneModel:
 			batches[updateOneCommand].models = append(batches[updateOneCommand].models, model)
+			batches[updateOneCommand].indexes = append(batches[updateOneCommand].indexes, i)
 		case *UpdateManyModel:
 			batches[updateManyCommand].models = append(batches[updateManyCommand].models, model)
+			batches[updateManyCommand].indexes = append(batches[updateManyCommand].indexes, i)
 		}
 	}
 
@@ -424,7 +426,7 @@ func createOrderedBatches(models []WriteModel) []bulkWriteBatch {
 	var prevKind writeCommandKind = -1
 	i := -1 // batch index
 
-	for _, model := range models {
+	for ind, model := range models {
 		var createNewBatch bool
 		var canRetry bool
 		var newKind writeCommandKind
@@ -455,13 +457,15 @@ func createOrderedBatches(models []WriteModel) []bulkWriteBatch {
 			batches = append(batches, bulkWriteBatch{
 				models:   []WriteModel{model},
 				canRetry: canRetry,
+				indexes:  []int{ind},
 			})
 			i++
 		} else {
 			batches[i].models = append(batches[i].models, model)
 			if !canRetry {
 				batches[i].canRetry = false // don't make it true if it was already false
 			}
+			batches[i].indexes = append(batches[i].indexes, ind)
 		}
 
 		prevKind = newKind
@@ -470,15 +474,15 @@ func createOrderedBatches(models []WriteModel) []bulkWriteBatch {
 	return batches
 }
 
-func (bw *bulkWrite) mergeResults(newResult BulkWriteResult, opIndex int64) {
+func (bw *bulkWrite) mergeResults(newResult BulkWriteResult) {
 	bw.result.InsertedCount += newResult.InsertedCount
 	bw.result.MatchedCount += newResult.MatchedCount
 	bw.result.ModifiedCount += newResult.ModifiedCount
 	bw.result.DeletedCount += newResult.DeletedCount
 	bw.result.UpsertedCount += newResult.UpsertedCount
 
 	for index, upsertID := range newResult.UpsertedIDs {
-		bw.result.UpsertedIDs[index+opIndex] = upsertID
+		bw.result.UpsertedIDs[index] = upsertID
 	}
 }
 

mongo/integration/collection_test.go

@@ -1067,6 +1067,77 @@ func TestCollection(t *testing.T) {
 			assert.Equal(mt, expectedModel, actualModel, "expected model %v in BulkWriteException, got %v",
 				expectedModel, actualModel)
 		})
+		mt.Run("unordered writeError index", func(mt *mtest.T) {
+			cappedOpts := bson.D{{"capped", true}, {"size", 64 * 1024}}
+			// Use a capped collection to get WriteErrors for delete operations
+			capped := mt.CreateCollection(mtest.Collection{
+				Name:       "deleteOne_capped",
+				CreateOpts: cappedOpts,
+			}, true)
+			models := []mongo.WriteModel{
+				mongo.NewInsertOneModel().SetDocument(bson.D{{"_id", "id1"}}),
+				mongo.NewInsertOneModel().SetDocument(bson.D{{"_id", "id3"}}),
+			}
+			_, err := capped.BulkWrite(mtest.Background, models, options.BulkWrite())
+			assert.Nil(t, err, "BulkWrite error: %v", err)
+
+			// UpdateOne and ReplaceOne models are batched together, so they each appear once
+			models = []mongo.WriteModel{
+				mongo.NewDeleteOneModel().SetFilter(bson.D{{"_id", "id0"}}),
+				mongo.NewDeleteManyModel().SetFilter(bson.D{{"_id", "id0"}}),
+				mongo.NewUpdateOneModel().SetFilter(bson.D{{"_id", "id3"}}).SetUpdate(bson.D{{"$set", bson.D{{"_id", 3.14159}}}}),
+				mongo.NewInsertOneModel().SetDocument(bson.D{{"_id", "id1"}}),
+				mongo.NewDeleteManyModel().SetFilter(bson.D{{"_id", "id0"}}),
+				mongo.NewUpdateManyModel().SetFilter(bson.D{{"_id", "id3"}}).SetUpdate(bson.D{{"$set", bson.D{{"_id", 3.14159}}}}),
+				mongo.NewDeleteOneModel().SetFilter(bson.D{{"_id", "id0"}}),
+				mongo.NewInsertOneModel().SetDocument(bson.D{{"_id", "id1"}}),
+				mongo.NewReplaceOneModel().SetFilter(bson.D{{"_id", "id3"}}).SetReplacement(bson.D{{"_id", 3.14159}}),
+				mongo.NewUpdateManyModel().SetFilter(bson.D{{"_id", "id3"}}).SetUpdate(bson.D{{"$set", bson.D{{"_id", 3.14159}}}}),
+			}
+			_, err = capped.BulkWrite(mtest.Background, models, options.BulkWrite().SetOrdered(false))
+			bwException, ok := err.(mongo.BulkWriteException)
+			assert.True(mt, ok, "expected error of type %T, got %T", mongo.BulkWriteException{}, err)
+
+			assert.Equal(mt, len(bwException.WriteErrors), 10, "expected 10 writeErrors, got %v", len(bwException.WriteErrors))
+			for _, writeErr := range bwException.WriteErrors {
+				switch writeErr.Request.(type) {
+				case *mongo.DeleteOneModel:
+					assert.True(mt, writeErr.Index == 0 || writeErr.Index == 6,
+						"expected index 0 or 6, got %v", writeErr.Index)
+				case *mongo.DeleteManyModel:
+					assert.True(mt, writeErr.Index == 1 || writeErr.Index == 4,
+						"expected index 1 or 4, got %v", writeErr.Index)
+				case *mongo.UpdateManyModel:
+					assert.True(mt, writeErr.Index == 5 || writeErr.Index == 9,
+						"expected index 5 or 9, got %v", writeErr.Index)
+				case *mongo.InsertOneModel:
+					assert.True(mt, writeErr.Index == 3 || writeErr.Index == 7,
+						"expected index 3 or 7, got %v", writeErr.Index)
+				case *mongo.UpdateOneModel:
+					assert.Equal(mt, writeErr.Index, 2, "expected index 2, got %v", writeErr.Index)
+				case *mongo.ReplaceOneModel:
+					assert.Equal(mt, writeErr.Index, 8, "expected index 8, got %v", writeErr.Index)
+				}
+
+			}
+		})
+		mt.Run("unordered upsertID index", func(mt *mtest.T) {
+			id1 := "id1"
+			id3 := "id3"
+			models := []mongo.WriteModel{
+				mongo.NewDeleteOneModel().SetFilter(bson.D{{"_id", "id0"}}),
+				mongo.NewReplaceOneModel().SetFilter(bson.D{{"_id", id1}}).SetReplacement(bson.D{{"_id", id1}}).SetUpsert(true),
+				mongo.NewDeleteOneModel().SetFilter(bson.D{{"_id", "id2"}}),
+				mongo.NewReplaceOneModel().SetFilter(bson.D{{"_id", id3}}).SetReplacement(bson.D{{"_id", id3}}).SetUpsert(true),
+				mongo.NewDeleteOneModel().SetFilter(bson.D{{"_id", "id4"}}),
+			}
+			res, err := mt.Coll.BulkWrite(mtest.Background, models, options.BulkWrite().SetOrdered(false))
+			assert.Nil(mt, err, "bulkwrite error: %v", err)
+
+			assert.Equal(mt, len(res.UpsertedIDs), 2, "expected 2 UpsertedIDs, got %v", len(res.UpsertedIDs))
+			assert.Equal(mt, res.UpsertedIDs[1].(string), id1, "expected UpsertedIDs[1] to be %v, got %v", id1, res.UpsertedIDs[1])
+			assert.Equal(mt, res.UpsertedIDs[3].(string), id3, "expected UpsertedIDs[3] to be %v, got %v", id3, res.UpsertedIDs[3])
+		})
 		unackClientOpts := options.Client().
 			SetWriteConcern(writeconcern.New(writeconcern.W(0)))
 		unackMtOpts := mtest.NewOptions().