diff --git a/collection.go b/collection.go index efe4689..ccf72cf 100644 --- a/collection.go +++ b/collection.go @@ -16,15 +16,17 @@ package qmgo import ( "context" "fmt" - "github.com/qiniu/qmgo/field" - "github.com/qiniu/qmgo/hook" - opts "github.com/qiniu/qmgo/options" + "reflect" + "strings" + "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" "go.mongodb.org/mongo-driver/x/bsonx" - "reflect" - "strings" + + "github.com/qiniu/qmgo/field" + "github.com/qiniu/qmgo/hook" + opts "github.com/qiniu/qmgo/options" ) // Collection is a handle to a MongoDB collection @@ -309,7 +311,6 @@ func (c *Collection) Aggregate(ctx context.Context, pipeline interface{}) Aggreg func (c *Collection) ensureIndex(ctx context.Context, indexes []opts.IndexModel) error { var indexModels []mongo.IndexModel - // 组建[]mongo.IndexModel for _, idx := range indexes { var model mongo.IndexModel var keysDoc bsonx.Doc @@ -375,17 +376,15 @@ func (c *Collection) EnsureIndexes(ctx context.Context, uniques []string, indexe // If the Key in opts.IndexModel is []string{"name"}, means create index: name // If the Key in opts.IndexModel is []string{"name","-age"} means create Compound indexes: name and -age func (c *Collection) CreateIndexes(ctx context.Context, indexes []opts.IndexModel) (err error) { - // 创建普通索引 err = c.ensureIndex(ctx, indexes) return } // CreateIndex creates one index // If the Key in opts.IndexModel is []string{"name"}, means create index name -// If the Key in opts.IndexModel is []string{"name","-age"} means drop Compound indexes: name and -age -func (c *Collection) CreateOneIndex(ctx context.Context, indexes opts.IndexModel) error { - // 创建普通索引 - return c.ensureIndex(ctx, []opts.IndexModel{indexes}) +// If the Key in opts.IndexModel is []string{"name","-age"} means create Compound index: name and -age +func (c *Collection) CreateOneIndex(ctx context.Context, index opts.IndexModel) error { + return c.ensureIndex(ctx, []opts.IndexModel{index}) } diff --git a/collection_test.go b/collection_test.go index a829e9d..4584783 100644 --- a/collection_test.go +++ b/collection_test.go @@ -15,13 +15,14 @@ package qmgo import ( "context" - "github.com/qiniu/qmgo/options" "testing" - "github.com/qiniu/qmgo/operator" "github.com/stretchr/testify/require" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" + + "github.com/qiniu/qmgo/operator" + "github.com/qiniu/qmgo/options" ) func TestCollection_EnsureIndex(t *testing.T) { diff --git a/cursor.go b/cursor.go index 174efeb..e29491e 100644 --- a/cursor.go +++ b/cursor.go @@ -15,6 +15,7 @@ package qmgo import ( "context" + "go.mongodb.org/mongo-driver/mongo" ) diff --git a/errors.go b/errors.go index 5e11ee4..ae185f5 100644 --- a/errors.go +++ b/errors.go @@ -41,6 +41,8 @@ var ( ErrNotSupportedPassword = errors.New("password not supported") // ErrNotValidSliceToInsert return if insert argument is not valid slice ErrNotValidSliceToInsert = errors.New("must be valid slice to insert") + // ErrReplacementContainUpdateOperators return if replacement document contain update operators + ErrReplacementContainUpdateOperators = errors.New("replacement document cannot contain keys beginning with '$'") ) // IsErrNoDocuments check if err is no documents, both mongo-go-driver error and qmgo custom error diff --git a/interface.go b/interface.go index a894aeb..aa9b544 100644 --- a/interface.go +++ b/interface.go @@ -28,12 +28,21 @@ package qmgo // EnsureIndexes(uniques []string, indexes []string) //} +// Change holds fields for running a findAndModify command via the Query.Apply method. +type Change struct { + Update interface{} // update/replace document + Replace bool // Whether to replace the document rather than updating + Remove bool // Whether to remove the document found rather than updating + Upsert bool // Whether to insert in case the document isn't found, take effect when Remove is false + ReturnNew bool // Should the modified document be returned rather than the old one, take effect when Remove is false +} + // CursorI Cursor interface type CursorI interface { Next(result interface{}) bool Close() error Err() error - All(reuslts interface{}) error + All(results interface{}) error //ID() int64 } @@ -48,6 +57,7 @@ type QueryI interface { Count() (n int64, err error) Distinct(key string, result interface{}) error Cursor() CursorI + Apply(change Change, result interface{}) error } // AggregateI define the interface of aggregate diff --git a/query.go b/query.go index 57202d4..e3bf93a 100644 --- a/query.go +++ b/query.go @@ -18,11 +18,12 @@ import ( "fmt" "reflect" - "github.com/qiniu/qmgo/hook" - qOpts "github.com/qiniu/qmgo/options" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" + + "github.com/qiniu/qmgo/hook" + qOpts "github.com/qiniu/qmgo/options" ) // Query struct definition @@ -274,3 +275,91 @@ func (q *Query) Cursor() CursorI { err: err, } } + +// Apply runs the findAndModify command, which allows updating, replacing +// or removing a document matching a query and atomically returning either the old +// version (the default) or the new version of the document (when ReturnNew is true) +// +// The Sort and Select query methods affect the result of Apply. In case +// multiple documents match the query, Sort enables selecting which document to +// act upon by ordering it first. Select enables retrieving only a selection +// of fields of the new or old document. +// +// When Change.Replace is true, it means replace at most one document in the collection +// and the update parameter must be a document and cannot contain any update operators; +// if no objects are found and Change.Upsert is false, it will returns ErrNoDocuments. +// When Change.Remove is true, it means delete at most one document in the collection +// and returns the document as it appeared before deletion; if no objects are found, +// it will returns ErrNoDocuments. +// When both Change.Replace and Change.Remove are false，it means update at most one document +// in the collection and the update parameter must be a document containing update operators; +// if no objects are found and Change.Upsert is false, it will returns ErrNoDocuments. +// +// reference: https://docs.mongodb.com/manual/reference/command/findAndModify/ +func (q *Query) Apply(change Change, result interface{}) error { + var err error + + if change.Remove { + err = q.findOneAndDelete(change, result) + } else if change.Replace { + err = q.findOneAndReplace(change, result) + } else { + err = q.findOneAndUpdate(change, result) + } + + return err +} + +// findOneAndDelete +// reference: https://docs.mongodb.com/manual/reference/method/db.collection.findOneAndDelete/ +func (q *Query) findOneAndDelete(change Change, result interface{}) error { + opts := options.FindOneAndDelete() + if q.sort != nil { + opts.SetSort(q.sort) + } + if q.project != nil { + opts.SetProjection(q.project) + } + + return q.collection.FindOneAndDelete(q.ctx, q.filter, opts).Decode(result) +} + +// findOneAndReplace +// reference: https://docs.mongodb.com/manual/reference/method/db.collection.findOneAndReplace/ +func (q *Query) findOneAndReplace(change Change, result interface{}) error { + opts := options.FindOneAndReplace() + if q.sort != nil { + opts.SetSort(q.sort) + } + if q.project != nil { + opts.SetProjection(q.project) + } + if change.Upsert { + opts.SetUpsert(change.Upsert) + } + if change.ReturnNew { + opts.SetReturnDocument(options.After) + } + + return q.collection.FindOneAndReplace(q.ctx, q.filter, change.Update, opts).Decode(result) +} + +// findOneAndUpdate +// reference: https://docs.mongodb.com/manual/reference/method/db.collection.findOneAndUpdate/ +func (q *Query) findOneAndUpdate(change Change, result interface{}) error { + opts := options.FindOneAndUpdate() + if q.sort != nil { + opts.SetSort(q.sort) + } + if q.project != nil { + opts.SetProjection(q.project) + } + if change.Upsert { + opts.SetUpsert(change.Upsert) + } + if change.ReturnNew { + opts.SetReturnDocument(options.After) + } + + return q.collection.FindOneAndUpdate(q.ctx, q.filter, change.Update, opts).Decode(result) +} \ No newline at end of file diff --git a/query_test.go b/query_test.go index d53dfd4..29aefb8 100644 --- a/query_test.go +++ b/query_test.go @@ -20,6 +20,9 @@ import ( "github.com/stretchr/testify/require" "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/bson/primitive" + "go.mongodb.org/mongo-driver/mongo" + + "github.com/qiniu/qmgo/operator" ) type QueryTestItem struct { @@ -563,3 +566,136 @@ func TestQuery_Cursor(t *testing.T) { cursor = cli.Find(context.Background(), filter3).Cursor() ast.Error(cursor.Err()) } + +func TestQuery_Apply(t *testing.T) { + ast := require.New(t) + cli := initClient("test") + defer cli.Close(context.Background()) + defer cli.DropCollection(context.Background()) + cli.EnsureIndexes(context.Background(), nil, []string{"name"}) + + id1 := primitive.NewObjectID() + id2 := primitive.NewObjectID() + id3 := primitive.NewObjectID() + docs := []interface{}{ + bson.M{"_id": id1, "name": "Alice", "age": 18}, + bson.M{"_id": id2, "name": "Alice", "age": 19}, + bson.M{"_id": id3, "name": "Lucas", "age": 20}, + } + _, _ = cli.InsertMany(context.Background(), docs) + + var err error + res1 := QueryTestItem{} + filter1 := bson.M{ + "name": "Tom", + } + change1 := Change{ + } + + err = cli.Find(context.Background(), filter1).Apply(change1, &res1) + ast.EqualError(err, mongo.ErrNilDocument.Error()) + + change1.Update = bson.M{ + operator.Set: bson.M{ + "name": "Tom", + "age": 18, + }, + } + err = cli.Find(context.Background(), filter1).Apply(change1, &res1) + ast.EqualError(err, mongo.ErrNoDocuments.Error()) + + change1.ReturnNew = true + err = cli.Find(context.Background(), filter1).Apply(change1, &res1) + ast.EqualError(err, mongo.ErrNoDocuments.Error()) + + change1.Upsert = true + err = cli.Find(context.Background(), filter1).Apply(change1, &res1) + ast.NoError(err) + ast.Equal( "Tom", res1.Name) + ast.Equal(18, res1.Age) + + res2 := QueryTestItem{} + filter2 := bson.M{ + "name": "Alice", + } + change2 := Change{ + ReturnNew: true, + Update: bson.M{ + operator.Set: bson.M{ + "name": "Alice", + "age": 22, + }, + }, + } + projection2 := bson.M{ + "age": 1, + } + err = cli.Find(context.Background(), filter2).Sort("age").Select(projection2).Apply(change2, &res2) + ast.NoError(err) + ast.Equal("", res2.Name) + ast.Equal(22, res2.Age) + + res3 := QueryTestItem{} + filter3 := bson.M{ + "name": "Bob", + } + change3 := Change{ + Remove: true, + } + err = cli.Find(context.Background(), filter3).Apply(change3, &res3) + ast.EqualError(err, mongo.ErrNoDocuments.Error()) + + res3 = QueryTestItem{} + filter3 = bson.M{ + "name": "Alice", + } + projection3 := bson.M{ + "age": 1, + } + err = cli.Find(context.Background(), filter3).Sort("age").Select(projection3).Apply(change3, &res3) + ast.NoError(err) + ast.Equal("", res3.Name) + ast.Equal(19, res3.Age) + + res4 := QueryTestItem{} + filter4 := bson.M{ + "name": "Bob", + } + change4 := Change{ + Replace: true, + Update: bson.M{ + operator.Set: bson.M{ + "name": "Bob", + "age": 23, + }, + }, + } + err = cli.Find(context.Background(), filter4).Apply(change4, &res4) + ast.EqualError(err, ErrReplacementContainUpdateOperators.Error()) + + change4.Update = bson.M{"name": "Bob", "age": 23} + err = cli.Find(context.Background(), filter4).Apply(change4, &res4) + ast.EqualError(err, mongo.ErrNoDocuments.Error()) + + change4.Upsert = true + change4.ReturnNew = true + err = cli.Find(context.Background(), filter4).Apply(change4, &res4) + ast.NoError(err) + ast.Equal("Bob", res4.Name) + ast.Equal(23, res4.Age) + + change4 = Change{ + Replace: true, + Update: bson.M{"name": "Bob", "age": 25}, + Upsert: true, + ReturnNew: false, + } + projection4 := bson.M{ + "age": 1, + "name": 1, + } + err = cli.Find(context.Background(), filter4).Sort("age").Select(projection4).Apply(change4, &res4) + ast.NoError(err) + ast.Equal("Bob", res4.Name) + ast.Equal(23, res4.Age) +} \ No newline at end of file diff --git a/util_test.go b/util_test.go index 0cd34ac..4533347 100644 --- a/util_test.go +++ b/util_test.go @@ -15,9 +15,10 @@ package qmgo import ( "fmt" - "github.com/stretchr/testify/require" "testing" "time" + + "github.com/stretchr/testify/require" ) func TestNow(t *testing.T) {