update:github.com/mongodb/mongo-go-driver v2.0 Migration (#4687)

2026-05-13 09:50:00 +08:00 · 2025-08-09 21:21:53 +08:00
parent f36e5fed35
commit b41b1b00df
22 changed files with 2515 additions and 1174 deletions
--- a/core/stores/mon/collection.go
+++ b/core/stores/mon/collection.go
@@ -1,3 +1,4 @@
+//go:generate mockgen -package mon -destination collection_mock.go -source collection.go Collection,monCollection
 package mon

 import (
@@ -8,9 +9,9 @@ import (
 	"github.com/zeromicro/go-zero/core/breaker"
 	"github.com/zeromicro/go-zero/core/errorx"
 	"github.com/zeromicro/go-zero/core/timex"
-	"go.mongodb.org/mongo-driver/mongo"
-	mopt "go.mongodb.org/mongo-driver/mongo/options"
-	"go.mongodb.org/mongo-driver/x/mongo/driver/session"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+	"go.mongodb.org/mongo-driver/v2/x/mongo/driver/session"
 )

 const (
@@ -47,79 +48,79 @@ type (
 	// Collection defines a MongoDB collection.
 	Collection interface {
 		// Aggregate executes an aggregation pipeline.
-		Aggregate(ctx context.Context, pipeline any, opts ...*mopt.AggregateOptions) (
+		Aggregate(ctx context.Context, pipeline any, opts ...options.Lister[options.AggregateOptions]) (
 			*mongo.Cursor, error)
 		// BulkWrite performs a bulk write operation.
-		BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...*mopt.BulkWriteOptions) (
+		BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...options.Lister[options.BulkWriteOptions]) (
 			*mongo.BulkWriteResult, error)
 		// Clone creates a copy of this collection with the same settings.
-		Clone(opts ...*mopt.CollectionOptions) (*mongo.Collection, error)
+		Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection
 		// CountDocuments returns the number of documents in the collection that match the filter.
-		CountDocuments(ctx context.Context, filter any, opts ...*mopt.CountOptions) (int64, error)
+		CountDocuments(ctx context.Context, filter any, opts ...options.Lister[options.CountOptions]) (int64, error)
 		// Database returns the database that this collection is a part of.
 		Database() *mongo.Database
 		// DeleteMany deletes documents from the collection that match the filter.
-		DeleteMany(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (
+		DeleteMany(ctx context.Context, filter any, opts ...options.Lister[options.DeleteManyOptions]) (
 			*mongo.DeleteResult, error)
 		// DeleteOne deletes at most one document from the collection that matches the filter.
-		DeleteOne(ctx context.Context, filter any, opts ...*mopt.DeleteOptions) (
+		DeleteOne(ctx context.Context, filter any, opts ...options.Lister[options.DeleteOneOptions]) (
 			*mongo.DeleteResult, error)
 		// Distinct returns a list of distinct values for the given key across the collection.
 		Distinct(ctx context.Context, fieldName string, filter any,
-			opts ...*mopt.DistinctOptions) ([]any, error)
+			opts ...options.Lister[options.DistinctOptions]) (*mongo.DistinctResult, error)
 		// Drop drops this collection from database.
-		Drop(ctx context.Context) error
+		Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error
 		// EstimatedDocumentCount returns an estimate of the count of documents in a collection
 		// using collection metadata.
-		EstimatedDocumentCount(ctx context.Context, opts ...*mopt.EstimatedDocumentCountOptions) (int64, error)
+		EstimatedDocumentCount(ctx context.Context, opts ...options.Lister[options.EstimatedDocumentCountOptions]) (int64, error)
 		// Find finds the documents matching the provided filter.
-		Find(ctx context.Context, filter any, opts ...*mopt.FindOptions) (*mongo.Cursor, error)
+		Find(ctx context.Context, filter any, opts ...options.Lister[options.FindOptions]) (*mongo.Cursor, error)
 		// FindOne returns up to one document that matches the provided filter.
-		FindOne(ctx context.Context, filter any, opts ...*mopt.FindOneOptions) (
+		FindOne(ctx context.Context, filter any, opts ...options.Lister[options.FindOneOptions]) (
 			*mongo.SingleResult, error)
 		// FindOneAndDelete returns at most one document that matches the filter. If the filter
 		// matches multiple documents, only the first document is deleted.
-		FindOneAndDelete(ctx context.Context, filter any, opts ...*mopt.FindOneAndDeleteOptions) (
+		FindOneAndDelete(ctx context.Context, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) (
 			*mongo.SingleResult, error)
 		// FindOneAndReplace returns at most one document that matches the filter. If the filter
 		// matches multiple documents, FindOneAndReplace returns the first document in the
 		// collection that matches the filter.
 		FindOneAndReplace(ctx context.Context, filter, replacement any,
-			opts ...*mopt.FindOneAndReplaceOptions) (*mongo.SingleResult, error)
+			opts ...options.Lister[options.FindOneAndReplaceOptions]) (*mongo.SingleResult, error)
 		// FindOneAndUpdate returns at most one document that matches the filter. If the filter
 		// matches multiple documents, FindOneAndUpdate returns the first document in the
 		// collection that matches the filter.
 		FindOneAndUpdate(ctx context.Context, filter, update any,
-			opts ...*mopt.FindOneAndUpdateOptions) (*mongo.SingleResult, error)
+			opts ...options.Lister[options.FindOneAndUpdateOptions]) (*mongo.SingleResult, error)
 		// Indexes returns the index view for this collection.
 		Indexes() mongo.IndexView
 		// InsertMany inserts the provided documents.
-		InsertMany(ctx context.Context, documents []any, opts ...*mopt.InsertManyOptions) (
+		InsertMany(ctx context.Context, documents []any, opts ...options.Lister[options.InsertManyOptions]) (
 			*mongo.InsertManyResult, error)
 		// InsertOne inserts the provided document.
-		InsertOne(ctx context.Context, document any, opts ...*mopt.InsertOneOptions) (
+		InsertOne(ctx context.Context, document any, opts ...options.Lister[options.InsertOneOptions]) (
 			*mongo.InsertOneResult, error)
 		// ReplaceOne replaces at most one document that matches the filter.
 		ReplaceOne(ctx context.Context, filter, replacement any,
-			opts ...*mopt.ReplaceOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error)
 		// UpdateByID updates a single document matching the provided filter.
 		UpdateByID(ctx context.Context, id, update any,
-			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
 		// UpdateMany updates the provided documents.
 		UpdateMany(ctx context.Context, filter, update any,
-			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error)
 		// UpdateOne updates a single document matching the provided filter.
 		UpdateOne(ctx context.Context, filter, update any,
-			opts ...*mopt.UpdateOptions) (*mongo.UpdateResult, error)
+			opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
 		// Watch returns a change stream cursor used to receive notifications of changes to the collection.
-		Watch(ctx context.Context, pipeline any, opts ...*mopt.ChangeStreamOptions) (
+		Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (
 			*mongo.ChangeStream, error)
 	}

 	decoratedCollection struct {
-		*mongo.Collection
-		name string
-		brk  breaker.Breaker
+		Collection monCollection
+		name       string
+		brk        breaker.Breaker
 	}

 	keepablePromise struct {
@@ -136,8 +137,16 @@ func newCollection(collection *mongo.Collection, brk breaker.Breaker) Collection
 	}
 }

+func newTestCollection(collection monCollection, brk breaker.Breaker) *decoratedCollection {
+	return &decoratedCollection{
+		Collection: collection,
+		name:       "test",
+		brk:        brk,
+	}
+}
+
 func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline any,
-	opts ...*mopt.AggregateOptions) (cur *mongo.Cursor, err error) {
+	opts ...options.Lister[options.AggregateOptions]) (cur *mongo.Cursor, err error) {
 	ctx, span := startSpan(ctx, aggregate)
 	defer func() {
 		endSpan(span, err)
@@ -157,7 +166,7 @@ func (c *decoratedCollection) Aggregate(ctx context.Context, pipeline any,
 }

 func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.WriteModel,
-	opts ...*mopt.BulkWriteOptions) (res *mongo.BulkWriteResult, err error) {
+	opts ...options.Lister[options.BulkWriteOptions]) (res *mongo.BulkWriteResult, err error) {
 	ctx, span := startSpan(ctx, bulkWrite)
 	defer func() {
 		endSpan(span, err)
@@ -177,7 +186,7 @@ func (c *decoratedCollection) BulkWrite(ctx context.Context, models []mongo.Writ
 }

 func (c *decoratedCollection) CountDocuments(ctx context.Context, filter any,
-	opts ...*mopt.CountOptions) (count int64, err error) {
+	opts ...options.Lister[options.CountOptions]) (count int64, err error) {
 	ctx, span := startSpan(ctx, countDocuments)
 	defer func() {
 		endSpan(span, err)
@@ -197,7 +206,7 @@ func (c *decoratedCollection) CountDocuments(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) DeleteMany(ctx context.Context, filter any,
-	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	opts ...options.Lister[options.DeleteManyOptions]) (res *mongo.DeleteResult, err error) {
 	ctx, span := startSpan(ctx, deleteMany)
 	defer func() {
 		endSpan(span, err)
@@ -217,7 +226,7 @@ func (c *decoratedCollection) DeleteMany(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) DeleteOne(ctx context.Context, filter any,
-	opts ...*mopt.DeleteOptions) (res *mongo.DeleteResult, err error) {
+	opts ...options.Lister[options.DeleteOneOptions]) (res *mongo.DeleteResult, err error) {
 	ctx, span := startSpan(ctx, deleteOne)
 	defer func() {
 		endSpan(span, err)
@@ -237,7 +246,7 @@ func (c *decoratedCollection) DeleteOne(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, filter any,
-	opts ...*mopt.DistinctOptions) (val []any, err error) {
+	opts ...options.Lister[options.DistinctOptions]) (res *mongo.DistinctResult, err error) {
 	ctx, span := startSpan(ctx, distinct)
 	defer func() {
 		endSpan(span, err)
@@ -249,7 +258,8 @@ func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, fi
 			c.logDurationSimple(ctx, distinct, startTime, err)
 		}()

-		val, err = c.Collection.Distinct(ctx, fieldName, filter, opts...)
+		res = c.Collection.Distinct(ctx, fieldName, filter, opts...)
+		err = res.Err()
 		return err
 	}, acceptable)

@@ -257,7 +267,7 @@ func (c *decoratedCollection) Distinct(ctx context.Context, fieldName string, fi
 }

 func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
-	opts ...*mopt.EstimatedDocumentCountOptions) (val int64, err error) {
+	opts ...options.Lister[options.EstimatedDocumentCountOptions]) (val int64, err error) {
 	ctx, span := startSpan(ctx, estimatedDocumentCount)
 	defer func() {
 		endSpan(span, err)
@@ -277,7 +287,7 @@ func (c *decoratedCollection) EstimatedDocumentCount(ctx context.Context,
 }

 func (c *decoratedCollection) Find(ctx context.Context, filter any,
-	opts ...*mopt.FindOptions) (cur *mongo.Cursor, err error) {
+	opts ...options.Lister[options.FindOptions]) (cur *mongo.Cursor, err error) {
 	ctx, span := startSpan(ctx, find)
 	defer func() {
 		endSpan(span, err)
@@ -297,7 +307,7 @@ func (c *decoratedCollection) Find(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) FindOne(ctx context.Context, filter any,
-	opts ...*mopt.FindOneOptions) (res *mongo.SingleResult, err error) {
+	opts ...options.Lister[options.FindOneOptions]) (res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOne)
 	defer func() {
 		endSpan(span, err)
@@ -318,7 +328,7 @@ func (c *decoratedCollection) FindOne(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter any,
-	opts ...*mopt.FindOneAndDeleteOptions) (res *mongo.SingleResult, err error) {
+	opts ...options.Lister[options.FindOneAndDeleteOptions]) (res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOneAndDelete)
 	defer func() {
 		endSpan(span, err)
@@ -339,7 +349,7 @@ func (c *decoratedCollection) FindOneAndDelete(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter any,
-	replacement any, opts ...*mopt.FindOneAndReplaceOptions) (
+	replacement any, opts ...options.Lister[options.FindOneAndReplaceOptions]) (
 	res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOneAndReplace)
 	defer func() {
@@ -361,7 +371,7 @@ func (c *decoratedCollection) FindOneAndReplace(ctx context.Context, filter any,
 }

 func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter, update any,
-	opts ...*mopt.FindOneAndUpdateOptions) (res *mongo.SingleResult, err error) {
+	opts ...options.Lister[options.FindOneAndUpdateOptions]) (res *mongo.SingleResult, err error) {
 	ctx, span := startSpan(ctx, findOneAndUpdate)
 	defer func() {
 		endSpan(span, err)
@@ -382,7 +392,7 @@ func (c *decoratedCollection) FindOneAndUpdate(ctx context.Context, filter, upda
 }

 func (c *decoratedCollection) InsertMany(ctx context.Context, documents []any,
-	opts ...*mopt.InsertManyOptions) (res *mongo.InsertManyResult, err error) {
+	opts ...options.Lister[options.InsertManyOptions]) (res *mongo.InsertManyResult, err error) {
 	ctx, span := startSpan(ctx, insertMany)
 	defer func() {
 		endSpan(span, err)
@@ -402,7 +412,7 @@ func (c *decoratedCollection) InsertMany(ctx context.Context, documents []any,
 }

 func (c *decoratedCollection) InsertOne(ctx context.Context, document any,
-	opts ...*mopt.InsertOneOptions) (res *mongo.InsertOneResult, err error) {
+	opts ...options.Lister[options.InsertOneOptions]) (res *mongo.InsertOneResult, err error) {
 	ctx, span := startSpan(ctx, insertOne)
 	defer func() {
 		endSpan(span, err)
@@ -422,7 +432,7 @@ func (c *decoratedCollection) InsertOne(ctx context.Context, document any,
 }

 func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter, replacement any,
-	opts ...*mopt.ReplaceOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.ReplaceOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, replaceOne)
 	defer func() {
 		endSpan(span, err)
@@ -442,7 +452,7 @@ func (c *decoratedCollection) ReplaceOne(ctx context.Context, filter, replacemen
 }

 func (c *decoratedCollection) UpdateByID(ctx context.Context, id, update any,
-	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, updateByID)
 	defer func() {
 		endSpan(span, err)
@@ -462,7 +472,7 @@ func (c *decoratedCollection) UpdateByID(ctx context.Context, id, update any,
 }

 func (c *decoratedCollection) UpdateMany(ctx context.Context, filter, update any,
-	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.UpdateManyOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, updateMany)
 	defer func() {
 		endSpan(span, err)
@@ -482,7 +492,7 @@ func (c *decoratedCollection) UpdateMany(ctx context.Context, filter, update any
 }

 func (c *decoratedCollection) UpdateOne(ctx context.Context, filter, update any,
-	opts ...*mopt.UpdateOptions) (res *mongo.UpdateResult, err error) {
+	opts ...options.Lister[options.UpdateOneOptions]) (res *mongo.UpdateResult, err error) {
 	ctx, span := startSpan(ctx, updateOne)
 	defer func() {
 		endSpan(span, err)
@@ -501,6 +511,27 @@ func (c *decoratedCollection) UpdateOne(ctx context.Context, filter, update any,
 	return
 }

+func (c *decoratedCollection) Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection {
+	return c.Collection.Clone(opts...)
+}
+
+func (c *decoratedCollection) Database() *mongo.Database {
+	return c.Collection.Database()
+}
+
+func (c *decoratedCollection) Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error {
+	return c.Collection.Drop(ctx, opts...)
+}
+
+func (c *decoratedCollection) Indexes() mongo.IndexView {
+	return c.Collection.Indexes()
+}
+
+func (c *decoratedCollection) Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (
+	*mongo.ChangeStream, error) {
+	return c.Collection.Watch(ctx, pipeline, opts...)
+}
+
 func (c *decoratedCollection) logDuration(ctx context.Context, method string,
 	startTime time.Duration, err error, docs ...any) {
 	logDurationWithDocs(ctx, c.name, method, startTime, err, docs...)
@@ -546,3 +577,73 @@ func isDupKeyError(err error) bool {

 	return e.HasErrorCode(duplicateKeyCode)
 }
+
+type (
+	// monCollection defines a MongoDB collection, used for unit test
+	monCollection interface {
+		// Aggregate executes an aggregation pipeline.
+		Aggregate(ctx context.Context, pipeline any, opts ...options.Lister[options.AggregateOptions]) (
+			*mongo.Cursor, error)
+		// BulkWrite performs a bulk write operation.
+		BulkWrite(ctx context.Context, models []mongo.WriteModel, opts ...options.Lister[options.BulkWriteOptions]) (
+			*mongo.BulkWriteResult, error)
+		// Clone creates a copy of this collection with the same settings.
+		Clone(opts ...options.Lister[options.CollectionOptions]) *mongo.Collection
+		// CountDocuments returns the number of documents in the collection that match the filter.
+		CountDocuments(ctx context.Context, filter any, opts ...options.Lister[options.CountOptions]) (int64, error)
+		// Database returns the database that this collection is a part of.
+		Database() *mongo.Database
+		// DeleteMany deletes documents from the collection that match the filter.
+		DeleteMany(ctx context.Context, filter any, opts ...options.Lister[options.DeleteManyOptions]) (
+			*mongo.DeleteResult, error)
+		// DeleteOne deletes at most one document from the collection that matches the filter.
+		DeleteOne(ctx context.Context, filter any, opts ...options.Lister[options.DeleteOneOptions]) (
+			*mongo.DeleteResult, error)
+		// Distinct returns a list of distinct values for the given key across the collection.
+		Distinct(ctx context.Context, fieldName string, filter any,
+			opts ...options.Lister[options.DistinctOptions]) *mongo.DistinctResult
+		// Drop drops this collection from database.
+		Drop(ctx context.Context, opts ...options.Lister[options.DropCollectionOptions]) error
+		// EstimatedDocumentCount returns an estimate of the count of documents in a collection
+		// using collection metadata.
+		EstimatedDocumentCount(ctx context.Context, opts ...options.Lister[options.EstimatedDocumentCountOptions]) (int64, error)
+		// Find finds the documents matching the provided filter.
+		Find(ctx context.Context, filter any, opts ...options.Lister[options.FindOptions]) (*mongo.Cursor, error)
+		// FindOne returns up to one document that matches the provided filter.
+		FindOne(ctx context.Context, filter any, opts ...options.Lister[options.FindOneOptions]) *mongo.SingleResult
+		// FindOneAndDelete returns at most one document that matches the filter. If the filter
+		// matches multiple documents, only the first document is deleted.
+		FindOneAndDelete(ctx context.Context, filter any, opts ...options.Lister[options.FindOneAndDeleteOptions]) *mongo.SingleResult
+		// FindOneAndReplace returns at most one document that matches the filter. If the filter
+		// matches multiple documents, FindOneAndReplace returns the first document in the
+		// collection that matches the filter.
+		FindOneAndReplace(ctx context.Context, filter, replacement any,
+			opts ...options.Lister[options.FindOneAndReplaceOptions]) *mongo.SingleResult
+		// FindOneAndUpdate returns at most one document that matches the filter. If the filter
+		// matches multiple documents, FindOneAndUpdate returns the first document in the
+		// collection that matches the filter.
+		FindOneAndUpdate(ctx context.Context, filter, update any,
+			opts ...options.Lister[options.FindOneAndUpdateOptions]) *mongo.SingleResult
+		// Indexes returns the index view for this collection.
+		Indexes() mongo.IndexView
+		// InsertMany inserts the provided documents.
+		InsertMany(ctx context.Context, documents interface{}, opts ...options.Lister[options.InsertManyOptions]) (*mongo.InsertManyResult, error)
+		// InsertOne inserts the provided document.
+		InsertOne(ctx context.Context, document any, opts ...options.Lister[options.InsertOneOptions]) (*mongo.InsertOneResult, error)
+		// ReplaceOne replaces at most one document that matches the filter.
+		ReplaceOne(ctx context.Context, filter, replacement any,
+			opts ...options.Lister[options.ReplaceOptions]) (*mongo.UpdateResult, error)
+		// UpdateByID updates a single document matching the provided filter.
+		UpdateByID(ctx context.Context, id, update any,
+			opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
+		// UpdateMany updates the provided documents.
+		UpdateMany(ctx context.Context, filter, update any,
+			opts ...options.Lister[options.UpdateManyOptions]) (*mongo.UpdateResult, error)
+		// UpdateOne updates a single document matching the provided filter.
+		UpdateOne(ctx context.Context, filter, update any,
+			opts ...options.Lister[options.UpdateOneOptions]) (*mongo.UpdateResult, error)
+		// Watch returns a change stream cursor used to receive notifications of changes to the collection.
+		Watch(ctx context.Context, pipeline any, opts ...options.Lister[options.ChangeStreamOptions]) (
+			*mongo.ChangeStream, error)
+	}
+)