-
Notifications
You must be signed in to change notification settings - Fork 0
/
mongo.go
88 lines (70 loc) · 1.77 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
80
81
82
83
84
85
86
87
88
package mongo
import (
"context"
"sync"
"time"
"go.mongodb.org/mongo-driver/mongo"
mongoOptions "go.mongodb.org/mongo-driver/mongo/options"
"go.mongodb.org/mongo-driver/mongo/readpref"
)
var (
pool = &sync.Map{} //实例池
mutex sync.Mutex //互斥锁
options = make(map[string]Option) //配置池
Default *mongo.Client //默认数据库
)
type Option struct {
DSN string `json:"dsn"` //数据源
}
func Add(name string, option Option) {
if option.DSN == "" {
panic("Option dsn empty " + name)
}
options[name] = option
}
func AddMap(name string, setting map[string]interface{}) {
option := Option{
DSN: setting["dsn"].(string),
}
Add(name, option)
}
func AddMapBatch(batch map[string]interface{}) {
for name, setting := range batch {
AddMap(name, setting.(map[string]interface{}))
}
}
func Use(name string) *mongo.Client {
if name == "" {
name = "default"
}
if instance, ok := pool.Load(name); ok {
return instance.(*mongo.Client)
} else {
mutex.Lock()
defer mutex.Unlock()
if instance, ok = pool.Load(name); ok {
return instance.(*mongo.Client)
}
}
option, ok := options[name]
if !ok {
panic("Option not found " + name)
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
newClient, err := mongo.Connect(ctx, mongoOptions.Client().ApplyURI(option.DSN))
if err != nil {
panic("Failed to connect mongo " + name + " err: " + err.Error())
}
ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
err = newClient.Ping(ctx, readpref.Primary())
if err != nil {
panic("Failed to connect mongo " + name + " err: " + err.Error())
}
if name == "default" {
Default = newClient
}
pool.Store(name, newClient)
return newClient
}