From 71ef4a26fe684323a0f15134f8cf66f0dd58bff3 Mon Sep 17 00:00:00 2001
From: xsean2020 <62974981+xsean2020@users.noreply.github.com>
Date: Fri, 20 Jan 2023 14:55:32 +0800
Subject: [PATCH] add array filter with apply (#267)

---
 collection.go |  1 +
 interface.go  |  1 +
 query.go      | 21 +++++++++++++++++++++
 query_test.go | 39 +++++++++++++++++++++++++++++++++++----
 4 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/collection.go b/collection.go
index cbd8425..76053ca 100644
--- a/collection.go
+++ b/collection.go
@@ -38,6 +38,7 @@ type Collection struct {
 
 // Find find by condition filter，return QueryI
 func (c *Collection) Find(ctx context.Context, filter interface{}, opts ...opts.FindOptions) QueryI {
+
 	return &Query{
 		ctx:        ctx,
 		collection: c.collection,
diff --git a/interface.go b/interface.go
index 61ed7e8..aef0d32 100644
--- a/interface.go
+++ b/interface.go
@@ -51,6 +51,7 @@ type CursorI interface {
 // QueryI Query interface
 type QueryI interface {
 	Collation(collation *options.Collation) QueryI
+	SetArrayFilters(*options.ArrayFilters) QueryI
 	Sort(fields ...string) QueryI
 	Select(selector interface{}) QueryI
 	Skip(n int64) QueryI
diff --git a/query.go b/query.go
index eb0833d..e90d356 100644
--- a/query.go
+++ b/query.go
@@ -33,6 +33,7 @@ type Query struct {
 	sort            interface{}
 	project         interface{}
 	hint            interface{}
+	arrayFilters    *options.ArrayFilters
 	limit           *int64
 	skip            *int64
 	batchSize       *int64
@@ -89,6 +90,22 @@ func (q *Query) Sort(fields ...string) QueryI {
 	return newQ
 }
 
+//  SetArrayFilter use for apply update array
+//  For Example :
+//  var res = QueryTestItem{}
+//  change := Change{
+//	Update:    bson.M{"$set": bson.M{"instock.$[elem].qty": 100}},
+//	ReturnNew: false,
+//  }
+//  cli.Find(context.Background(), bson.M{"name": "Lucas"}).
+//      SetArrayFilters(&options.ArrayFilters{Filters: []interface{}{bson.M{"elem.warehouse": bson.M{"$in": []string{"C", "F"}}},}}).
+//        Apply(change, &res)
+func (q *Query) SetArrayFilters(filter *options.ArrayFilters) QueryI {
+	newQ := q
+	newQ.arrayFilters = filter
+	return newQ
+}
+
 // Select is used to determine which fields are displayed or not displayed in the returned results
 // Format: bson.M{"age": 1} means that only the age field is displayed
 // bson.M{"age": 0} means to display other fields except age
@@ -405,6 +422,10 @@ func (q *Query) findOneAndUpdate(change Change, result interface{}) error {
 		opts.SetReturnDocument(options.After)
 	}
 
+	if q.arrayFilters != nil {
+		opts.SetArrayFilters(*q.arrayFilters)
+	}
+
 	err := q.collection.FindOneAndUpdate(q.ctx, q.filter, change.Update, opts).Decode(result)
 	if change.Upsert && !change.ReturnNew && err == mongo.ErrNoDocuments {
 		return nil
diff --git a/query_test.go b/query_test.go
index ff817b1..72e6d7e 100644
--- a/query_test.go
+++ b/query_test.go
@@ -19,18 +19,23 @@ import (
 	"testing"
 	"time"
 
+	"github.com/qiniu/qmgo/operator"
 	"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"
+	"go.mongodb.org/mongo-driver/mongo/options"
 )
 
 type QueryTestItem struct {
 	Id   primitive.ObjectID `bson:"_id"`
 	Name string             `bson:"name"`
 	Age  int                `bson:"age"`
+
+	Instock []struct {
+		Warehouse string `bson:"warehouse"`
+		Qty       int    `bson:"qty"`
+	} `bson:"instock"`
 }
 
 type QueryTestItem2 struct {
@@ -679,8 +684,13 @@ func TestQuery_Apply(t *testing.T) {
 	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},
-	}
+		bson.M{"_id": id3, "name": "Lucas", "age": 20, "instock": []bson.M{
+			{"warehouse": "B", "qty": 15},
+			{"warehouse": "C", "qty": 35},
+			{"warehouse": "E", "qty": 15},
+			{"warehouse": "F", "qty": 45},
+		}}}
+
 	_, _ = cli.InsertMany(context.Background(), docs)
 
 	var err error
@@ -829,6 +839,26 @@ func TestQuery_Apply(t *testing.T) {
 	ast.NoError(err)
 	ast.Equal("", res4.Name)
 	ast.Equal(0, res4.Age)
+
+	var res5 = QueryTestItem{}
+	filter5 := bson.M{"name": "Lucas"}
+	change5 := Change{
+		Update:    bson.M{"$set": bson.M{"instock.$[elem].qty": 100}},
+		ReturnNew: true,
+	}
+	err = cli.Find(context.Background(), filter5).SetArrayFilters(&options.ArrayFilters{Filters: []interface{}{
+		bson.M{"elem.warehouse": bson.M{"$in": []string{"C", "F"}}},
+	}}).Apply(change5, &res5)
+	ast.NoError(err)
+
+	for _, item := range res5.Instock {
+		switch item.Warehouse {
+		case "C", "F":
+			ast.Equal(100, item.Qty)
+		case "B", "E":
+			ast.Equal(15, item.Qty)
+		}
+	}
 }
 
 func TestQuery_BatchSize(t *testing.T) {
@@ -854,4 +884,5 @@ func TestQuery_BatchSize(t *testing.T) {
 	err := cli.Find(context.Background(), bson.M{"name": "Alice"}).BatchSize(1).All(&res)
 	ast.NoError(err)
 	ast.Len(res, 2)
+
 }