/
mongo_connection_manager.go
71 lines (59 loc) · 1.96 KB
/
mongo_connection_manager.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
package mongodb_storage
import (
"context"
"github.com/storage-lock/go-storage"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
"sync"
)
// MongoConnectionManager 负责维护与Mongo数据库的连接
type MongoConnectionManager struct {
// 连接到数据库的地址
URI string
// Mongo客户端
clientOnce sync.Once
err error
client *mongo.Client
}
var _ storage.ConnectionManager[*mongo.Client] = &MongoConnectionManager{}
// NewMongoConnectionManagerFromURI 从Mongo uri创建连接管理器
func NewMongoConnectionManagerFromURI(uri string) *MongoConnectionManager {
return &MongoConnectionManager{
URI: uri,
}
}
// NewMongoConnectionManagerFromClient 复用已经存在的mongo client,从其创建连接管理器
func NewMongoConnectionManagerFromClient(client *mongo.Client) *MongoConnectionManager {
return &MongoConnectionManager{
client: client,
}
}
const MongoConnectionManagerName = "mongodb-connection-manager"
func (x *MongoConnectionManager) Name() string {
return MongoConnectionManagerName
}
func (x *MongoConnectionManager) Take(ctx context.Context) (*mongo.Client, error) {
x.clientOnce.Do(func() {
// 当前还没有创建连接,并且也没有过错误(err主要是用来标记之前连接时发生的错误,有的话就不重复调用了)
if x.client == nil && x.err == nil {
// 此处暂不支持定制,如果想修改连接参数的话可以自行创建client从client创建ConnectionManager
client, err := mongo.Connect(ctx, options.Client().ApplyURI(x.URI))
if err != nil {
x.err = err
return
}
x.client = client
}
})
return x.client, x.err
}
func (x *MongoConnectionManager) Return(ctx context.Context, connection *mongo.Client) error {
return nil
}
func (x *MongoConnectionManager) Shutdown(ctx context.Context) error {
// 连接池关闭的时候断掉连接,即使它是被复用的
if x.client != nil {
return x.client.Disconnect(ctx)
}
return nil
}