Add custom options to client bulkWrite.

Author: qingyang-hu

internal/integration/client_test.go

@@ -31,7 +31,9 @@ import (
 	"go.mongodb.org/mongo-driver/v2/mongo/writeconcern"
 	"go.mongodb.org/mongo-driver/v2/x/bsonx/bsoncore"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/description"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/wiremessage"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/xoptions"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -838,8 +840,40 @@ func TestClient_BulkWrite(t *testing.T) {
 		}
 
 		_, err := mt.Client.BulkWrite(context.Background(), writes)
-		require.NoError(t, err)
-		assert.Equal(t, 2, bulkWrites, "expected %d bulkWrites, got %d", 2, bulkWrites)
+		require.NoError(mt, err)
+		assert.Equal(mt, 2, bulkWrites, "expected %d bulkWrites, got %d", 2, bulkWrites)
+	})
+	mt.Run("test options callback", func(mt *mtest.T) {
+		mt.Parallel()
+
+		insertOneModel := mongo.NewClientInsertOneModel().SetDocument(bson.D{{"x", 1}})
+		writes := []mongo.ClientBulkWrite{{
+			Database:   "foo",
+			Collection: "bar",
+			Model:      insertOneModel,
+		}}
+
+		marshalValue := func(val interface{}) bson.RawValue {
+			t.Helper()
+
+			valType, data, err := bson.MarshalValue(val)
+			require.Nil(t, err, "MarshalValue error: %v", err)
+			return bson.RawValue{
+				Type:  valType,
+				Value: data,
+			}
+		}
+
+		opts := options.ClientBulkWrite()
+		xoptions.SetInternalClientBulkWriteOptions(opts, "commandCallback", func(dst []byte, _ description.SelectedServer) ([]byte, error) {
+			dst = bsoncore.AppendStringElement(dst, "foo", "bar")
+			return dst, nil
+		})
+		_, _ = mt.Client.BulkWrite(context.Background(), writes, opts)
+		evt := mt.GetStartedEvent()
+		val := evt.Command.Lookup("foo")
+		expected := marshalValue("bar")
+		assert.Equal(mt, expected, val, "expected value to be %s", expected.String())
 	})
 }
 

mongo/client.go

@@ -963,6 +963,16 @@ func (c *Client) BulkWrite(ctx context.Context, writes []ClientBulkWrite,
 			op.rawData = &rawData
 		}
 	}
+	if bypassEmptyTsReplacementOpt := optionsutil.Value(bwo.Internal, "bypassEmptyTsReplacement"); bypassEmptyTsReplacementOpt != nil {
+		if bypassEmptyTsReplacement, ok := bypassEmptyTsReplacementOpt.(bool); ok {
+			op.bypassEmptyTsReplacement = &bypassEmptyTsReplacement
+		}
+	}
+	if commandCallbackOpt := optionsutil.Value(bwo.Internal, "commandCallback"); commandCallbackOpt != nil {
+		if commandCallback, ok := commandCallbackOpt.(func([]byte, description.SelectedServer) ([]byte, error)); ok {
+			op.commandCallback = commandCallback
+		}
+	}
 	if bwo.VerboseResults == nil || !(*bwo.VerboseResults) {
 		op.errorsOnly = true
 	} else if !acknowledged {

mongo/client_bulk_write.go

@@ -45,6 +45,8 @@ type clientBulkWrite struct {
 	selector                 description.ServerSelector
 	writeConcern             *writeconcern.WriteConcern
 	rawData                  *bool
+	bypassEmptyTsReplacement *bool
+	commandCallback          func([]byte, description.SelectedServer) ([]byte, error)
 
 	result ClientBulkWriteResult
 }
@@ -122,7 +124,8 @@ func (bw *clientBulkWrite) execute(ctx context.Context) error {
 }
 
 func (bw *clientBulkWrite) newCommand() func([]byte, description.SelectedServer) ([]byte, error) {
-	return func(dst []byte, desc description.SelectedServer) ([]byte, error) {
+	return func(cmd []byte, desc description.SelectedServer) ([]byte, error) {
+		var dst []byte
 		dst = bsoncore.AppendInt32Element(dst, "bulkWrite", 1)
 
 		dst = bsoncore.AppendBooleanElement(dst, "errorsOnly", bw.errorsOnly)
@@ -148,7 +151,19 @@ func (bw *clientBulkWrite) newCommand() func([]byte, description.SelectedServer)
 		if bw.rawData != nil && desc.WireVersion != nil && driverutil.VersionRangeIncludes(*desc.WireVersion, 27) {
 			dst = bsoncore.AppendBooleanElement(dst, "rawData", *bw.rawData)
 		}
-		return dst, nil
+		if bw.bypassEmptyTsReplacement != nil {
+			dst = bsoncore.AppendBooleanElement(dst, "bypassEmptyTsReplacement", *bw.bypassEmptyTsReplacement)
+		}
+		if bw.commandCallback != nil {
+			var err error
+			dst, err = bw.commandCallback(dst, desc)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		cmd = append(cmd, dst...)
+		return cmd, nil
 	}
 }
 

x/mongo/driver/xoptions/options.go

@@ -12,6 +12,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/internal/optionsutil"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/x/mongo/driver"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/description"
 )
 
 // SetInternalClientOptions sets internal options for ClientOptions.
@@ -101,6 +102,24 @@ func SetInternalClientBulkWriteOptions(a *options.ClientBulkWriteOptionsBuilder,
 			opts.Internal = optionsutil.WithValue(opts.Internal, key, b)
 			return nil
 		})
+	case "bypassEmptyTsReplacement":
+		b, ok := option.(bool)
+		if !ok {
+			return typeErrFunc("bool")
+		}
+		a.Opts = append(a.Opts, func(opts *options.ClientBulkWriteOptions) error {
+			opts.Internal = optionsutil.WithValue(opts.Internal, key, b)
+			return nil
+		})
+	case "commandCallback":
+		cb, ok := option.(func([]byte, description.SelectedServer) ([]byte, error))
+		if !ok {
+			return typeErrFunc("func([]byte, description.SelectedServer) ([]byte, error)")
+		}
+		a.Opts = append(a.Opts, func(opts *options.ClientBulkWriteOptions) error {
+			opts.Internal = optionsutil.WithValue(opts.Internal, key, cb)
+			return nil
+		})
 	default:
 		return fmt.Errorf("unsupported option: %q", key)
 	}