WIP

Author: qingyang-hu

mongo/client_bulk_write.go

@@ -58,6 +58,7 @@ func (bw *clientBulkWrite) execute(ctx context.Context) error {
 	}
 	resMap := make([]interface{}, len(bw.models))
 	insIDMap := make(map[int]interface{})
+	canRetry := true
 	for i, v := range bw.models {
 		var doc bsoncore.Document
 		var err error
@@ -94,6 +95,7 @@ func (bw *clientBulkWrite) execute(ctx context.Context) error {
 				bw.client.bsonOpts,
 				bw.client.registry)
 		case *ClientUpdateManyModel:
+			canRetry = false
 			nsIdx = getNsIndex(model.Namespace)
 			if bw.result.UpdateResults == nil {
 				bw.result.UpdateResults = make(map[int64]ClientUpdateResult)
@@ -144,6 +146,7 @@ func (bw *clientBulkWrite) execute(ctx context.Context) error {
 				bw.client.bsonOpts,
 				bw.client.registry)
 		case *ClientDeleteManyModel:
+			canRetry = false
 			nsIdx = getNsIndex(model.Namespace)
 			if bw.result.DeleteResults == nil {
 				bw.result.DeleteResults = make(map[int64]ClientDeleteResult)
@@ -168,24 +171,22 @@ func (bw *clientBulkWrite) execute(ctx context.Context) error {
 		Documents:  docs,
 		Ordered:    bw.ordered,
 	}
+	retry := driver.RetryNone
+	if bw.client.retryWrites && canRetry {
+		retry = driver.RetryOncePerCommand
+	}
 	op := operation.NewCommandFn(bw.newCommand(nsList)).Batches(batches).
 		Session(bw.session).WriteConcern(bw.writeConcern).CommandMonitor(bw.client.monitor).
 		ServerSelector(bw.selector).ClusterClock(bw.client.clock).
-		Database("admin").
+		Database("admin").Type(driver.Write).RetryMode(retry).
 		Deployment(bw.client.deployment).Crypt(bw.client.cryptFLE).
 		ServerAPI(bw.client.serverAPI).Timeout(bw.client.timeout).
 		Logger(bw.client.logger).Authenticator(bw.client.authenticator).Name("bulkWrite")
-	opErr := op.Execute(ctx)
-	var wcErrs []*WriteConcernError
-	if opErr != nil {
-		if errors.Is(opErr, driver.ErrUnacknowledgedWrite) {
+	err := op.Execute(ctx)
+	if err != nil {
+		if errors.Is(err, driver.ErrUnacknowledgedWrite) {
 			return nil
 		}
-		var writeErr driver.WriteCommandError
-		if errors.As(opErr, &writeErr) {
-			wcErr := convertDriverWriteConcernError(writeErr.WriteConcernError)
-			wcErrs = append(wcErrs, wcErr)
-		}
 	}
 	var res struct {
 		Ok     bool
@@ -212,34 +213,33 @@ func (bw *clientBulkWrite) execute(ctx context.Context) error {
 	bw.result.UpsertedCount = int64(res.NUpserted)
 	errors := make(map[int64]WriteError)
 	for i, cur := range res.Cursor.FirstBatch {
-		switch res := resMap[i].(type) {
+		switch r := resMap[i].(type) {
 		case map[int64]ClientDeleteResult:
-			if err := appendDeleteResult(cur, res, errors); err != nil {
+			if err := appendDeleteResult(cur, r, errors); err != nil {
 				return err
 			}
 		case map[int64]ClientInsertResult:
-			if err := appendInsertResult(cur, res, errors, insIDMap); err != nil {
+			if err := appendInsertResult(cur, r, errors, insIDMap); err != nil {
 				return err
 			}
 		case map[int64]ClientUpdateResult:
-			if err := appendUpdateResult(cur, res, errors); err != nil {
+			if err := appendUpdateResult(cur, r, errors); err != nil {
 				return err
 			}
 		}
 	}
-	if !res.Ok || res.NErrors > 0 || opErr != nil {
+	if !res.Ok || res.NErrors > 0 {
 		return ClientBulkWriteException{
 			TopLevelError: &WriteError{
 				Code:    int(res.Code),
 				Message: res.Errmsg,
 				Raw:     bson.Raw(rawRes),
 			},
-			WriteConcernErrors: wcErrs,
-			WriteErrors:        errors,
-			PartialResult:      &bw.result,
+			WriteErrors:   errors,
+			PartialResult: &bw.result,
 		}
 	}
-	return nil
+	return err
 }
 
 func (bw *clientBulkWrite) newCommand(nsList []string) func([]byte, description.SelectedServer) ([]byte, error) {

x/mongo/driver/operation.go

@@ -598,7 +598,7 @@ func (op Operation) Execute(ctx context.Context) error {
 		if conn != nil {
 			// If we are dealing with a sharded cluster, then mark the failed server
 			// as "deprioritized".
-			if desc := conn.Description; desc != nil && op.Deployment.Kind() == description.Sharded {
+			if op.Deployment.Kind() == description.Sharded {
 				deprioritizedServers = []description.Server{conn.Description()}
 			}
 

x/mongo/driver/operation/command.go

@@ -26,13 +26,15 @@ type Command struct {
 	name           string
 	authenticator  driver.Authenticator
 	commandFn      func([]byte, description.SelectedServer) ([]byte, error)
-	batches        *driver.Batches
 	database       string
 	deployment     driver.Deployment
 	selector       description.ServerSelector
 	writeConcern   *writeconcern.WriteConcern
 	readPreference *readpref.ReadPref
 	clock          *session.ClusterClock
+	retry          *driver.RetryMode
+	opType         driver.Type
+	batches        *driver.Batches
 	session        *session.Client
 	monitor        *event.CommandMonitor
 	resultResponse bsoncore.Document
@@ -95,7 +97,6 @@ func (c *Command) Execute(ctx context.Context) error {
 
 	return driver.Operation{
 		CommandFn: c.commandFn,
-		Batches:   c.batches,
 		ProcessResponseFn: func(info driver.ResponseInfo) error {
 			c.resultResponse = info.ServerResponse
 
@@ -113,6 +114,9 @@ func (c *Command) Execute(ctx context.Context) error {
 		},
 		Client:         c.session,
 		Clock:          c.clock,
+		RetryMode:      c.retry,
+		Type:           c.opType,
+		Batches:        c.batches,
 		CommandMonitor: c.monitor,
 		Database:       c.database,
 		Deployment:     c.deployment,
@@ -168,6 +172,26 @@ func (c *Command) Batches(batches *driver.Batches) *Command {
 	return c
 }
 
+// RetryMode sets the RetryMode for this operation.
+func (c *Command) RetryMode(retry driver.RetryMode) *Command {
+	if c == nil {
+		c = new(Command)
+	}
+
+	c.retry = &retry
+	return c
+}
+
+// Type sets the opType for this operation.
+func (c *Command) Type(t driver.Type) *Command {
+	if c == nil {
+		c = new(Command)
+	}
+
+	c.opType = t
+	return c
+}
+
 // Database sets the database to run this operation against.
 func (c *Command) Database(database string) *Command {
 	if c == nil {