forked from zeromicro/go-zero
-
Notifications
You must be signed in to change notification settings - Fork 0
/
clientmanager.go
57 lines (45 loc) · 1.33 KB
/
clientmanager.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
package mon
import (
"context"
"io"
"time"
"github.com/zeromicro/go-zero/core/syncx"
"go.mongodb.org/mongo-driver/mongo"
mopt "go.mongodb.org/mongo-driver/mongo/options"
)
const defaultTimeout = time.Second
var clientManager = syncx.NewResourceManager()
// ClosableClient wraps *mongo.Client and provides a Close method.
type ClosableClient struct {
*mongo.Client
}
// Close disconnects the underlying *mongo.Client.
func (cs *ClosableClient) Close() error {
return cs.Client.Disconnect(context.Background())
}
// Inject injects a *mongo.Client into the client manager.
// Typically, this is used to inject a *mongo.Client for test purpose.
func Inject(key string, client *mongo.Client) {
clientManager.Inject(key, &ClosableClient{client})
}
func getClient(url string, clientOpts ...*mopt.ClientOptions) (*mongo.Client, error) {
val, err := clientManager.GetResource(url, func() (io.Closer, error) {
opt := mopt.Client().ApplyURI(url)
if len(clientOpts) > 0 {
// If options has conflict, opt from URL win.
opt = mopt.MergeClientOptions(clientOpts[0], opt)
}
cli, err := mongo.Connect(context.Background(), opt)
if err != nil {
return nil, err
}
concurrentSess := &ClosableClient{
Client: cli,
}
return concurrentSess, nil
})
if err != nil {
return nil, err
}
return val.(*ClosableClient).Client, nil
}