/
common.go
125 lines (98 loc) · 2.85 KB
/
common.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
package mongodb
import (
"context"
"reflect"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
// SelectOne 通过反射查询单条记录
func SelectOne(col *mongo.Collection, filter interface{}, model reflect.Type, options ...*options.FindOneOptions) (interface{}, error) {
var res = Elem(model).Addr().Interface()
if err := FindOne(col, filter, res, options...); err != nil {
return nil, err
}
return res, nil
}
// SelectAll 通过反射查询多条记录
func SelectAll(col *mongo.Collection, filter interface{}, model reflect.Type, options ...*options.FindOptions) ([]interface{}, error) {
rows := SliceElem(model)
if err := FindAll(col, filter, rows.Addr().Interface(), options...); err != nil {
return nil, err
}
if rows.IsNil() {
return nil, nil
}
var result []interface{}
for i := 0; i < rows.Len(); i++ {
result = append(result, rows.Index(i).Interface())
}
return result, nil
}
func FindOne(col *mongo.Collection, filter interface{}, result interface{}, options ...*options.FindOneOptions) error {
ctx, cancel := context.WithTimeout(context.Background(), DefaultReadWriteTimeout)
defer cancel()
if filter == nil {
filter = bson.M{}
}
if err := col.FindOne(ctx, filter, options...).Decode(result); err != nil {
return err
}
return nil
}
func FindAll(col *mongo.Collection, filter interface{}, result interface{}, options ...*options.FindOptions) error {
ctx, cancel := context.WithTimeout(context.Background(), DefaultReadWriteTimeout)
defer cancel()
if filter == nil {
filter = bson.M{}
}
cur, err := col.Find(ctx, filter, options...)
if err != nil {
return err
}
defer cur.Close(ctx)
if err := cur.All(context.Background(), result); err != nil {
return err
}
return nil
}
// 分段获取数据
func FindScan(ctx context.Context, col *mongo.Collection, page, size int64, filter interface{}, result interface{}, fn ...func(opts *options.FindOptions) *options.FindOptions) *Scan {
if filter == nil {
filter = bson.M{}
}
var scan = new(Scan)
count, err := col.CountDocuments(ctx, filter)
if err != nil {
return scan
}
scan = NewScan(count, page, size)
if count > 0 {
opts := scan.FindOptions()
if len(fn) > 0 && fn[0] != nil {
opts = fn[0](opts)
}
cur, err := col.Find(ctx, filter, opts)
if err != nil {
return scan
}
defer cur.Close(ctx)
if err := cur.All(ctx, result); err != nil {
return scan
}
}
return scan
}
func Aggregate(col *mongo.Collection, pipeline interface{}, result interface{}, opts ...*options.AggregateOptions) error {
ctx, cancel := context.WithTimeout(context.Background(), DefaultReadWriteTimeout)
defer cancel()
if pipeline == nil {
pipeline = bson.M{}
}
cur, err := col.Aggregate(ctx, pipeline, opts...)
if err != nil {
return err
}
defer cur.Close(ctx)
return cur.All(ctx, result)
}