/
mongo.go
79 lines (68 loc) · 1.59 KB
/
mongo.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
package mongo
import (
"context"
"time"
"github.com/ops-cn/admin/app/model/impl/mongo/entity"
"github.com/ops-cn/common/logger"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
// Config 配置参数
type Config struct {
URI string
Database string
Timeout time.Duration
}
// NewClient 创建mongo客户端实例
func NewClient(cfg *Config) (*mongo.Client, func(), error) {
var (
ctx = context.Background()
cancel context.CancelFunc
)
if t := cfg.Timeout; t > 0 {
ctx, cancel = context.WithTimeout(ctx, t)
defer cancel()
}
cli, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.URI))
if err != nil {
return nil, nil, err
}
cleanFunc := func() {
err := cli.Disconnect(context.Background())
if err != nil {
logger.Errorf(context.Background(), "Mongo disconnect error: %s", err.Error())
}
}
err = cli.Ping(context.Background(), nil)
if err != nil {
return nil, cleanFunc, err
}
return cli, cleanFunc, nil
}
// CreateIndexes 创建索引
func CreateIndexes(ctx context.Context, cli *mongo.Client) error {
return createIndexes(
ctx,
cli,
new(entity.Demo),
new(entity.MenuAction),
new(entity.MenuActionResource),
new(entity.Menu),
new(entity.RoleMenu),
new(entity.Role),
new(entity.UserRole),
new(entity.User),
)
}
type indexer interface {
CreateIndexes(ctx context.Context, cli *mongo.Client) error
}
func createIndexes(ctx context.Context, cli *mongo.Client, indexes ...indexer) error {
for _, idx := range indexes {
err := idx.CreateIndexes(ctx, cli)
if err != nil {
return err
}
}
return nil
}