Merge pull request #95 from qnxm/develop

add Apply func

Author: jiangz222

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})
 
 }
 

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) {

cursor.go

@@ -15,6 +15,7 @@ package qmgo
 
 import (
 	"context"
+
 	"go.mongodb.org/mongo-driver/mongo"
 )
 

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

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

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)
+}

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)
+}

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) {