/
client.go
144 lines (122 loc) · 4.17 KB
/
client.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
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
/*
Copyright 2021 The KodeRover Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package mongo
import (
"context"
"reflect"
"sync"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/bsoncodec"
"go.mongodb.org/mongo-driver/bson/bsonoptions"
"go.mongodb.org/mongo-driver/bson/bsontype"
"go.mongodb.org/mongo-driver/mongo/readpref"
"github.com/koderover/zadig/v2/pkg/microservice/aslan/config"
"github.com/koderover/zadig/v2/pkg/tool/log"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
var once sync.Once
var client *mongo.Client
func Database(name string) *mongo.Database {
return Client().Database(name)
}
func SessionContext(ctx context.Context, session mongo.Session) context.Context {
if !config.EnableTransaction() {
return ctx
}
if session == nil {
return ctx
}
return mongo.NewSessionContext(ctx, session)
}
func Session() mongo.Session {
session, err := Client().StartSession()
if err != nil {
log.Panicf("Failed to start mongo session, err: %v", err)
return nil
}
return session
}
func StartTransaction(session mongo.Session) error {
if !config.EnableTransaction() {
return nil
}
return session.StartTransaction()
}
func AbortTransaction(session mongo.Session) error {
if !config.EnableTransaction() {
return nil
}
return session.AbortTransaction(context.TODO())
}
func CommitTransaction(session mongo.Session) error {
if !config.EnableTransaction() {
return nil
}
return session.CommitTransaction(context.Background())
}
func Client() *mongo.Client {
if client == nil {
panic("mongoDB connection is not initialized yet")
}
return client
}
// Init is a singleton, it will be initialized only once.
// In case the uri provides only a single host in the mongodb cluster, the system will
// attempt to connect without discovering other hosts in the cluster.
func Init(ctx context.Context, uri string) {
once.Do(func() {
nilSliceCodec := bsoncodec.NewSliceCodec(bsonoptions.SliceCodec().SetEncodeNilAsEmpty(true))
tM := reflect.TypeOf(bson.M{})
reg := bson.NewRegistryBuilder().RegisterTypeMapEntry(bsontype.EmbeddedDocument, tM).RegisterDefaultEncoder(reflect.Slice, nilSliceCodec).Build()
connInfo, err := extractURL(uri)
if err != nil {
log.Fatalf("Failed to initialize mongo db connection, err: %v", err)
}
opt := options.Client().ApplyURI(uri).SetRegistry(reg)
// By default the client will discover the mongodb cluster topology (if exists) and try to
// connect to ALL hosts in the cluster.
// If NONE of the host is discoverable by its host name (private network host name),
// and only a single host ip is provided, the auto-discovery function will cause a panic due
// to non of the host can be connected by the discovered host name.
// Thus, when there is only 1 addr in the provided uri, the system will try to connect with
// the given connection string ONLY.
// ref: https://pkg.go.dev/go.mongodb.org/mongo-driver/mongo/options#ClientOptions.SetDirect
if len(connInfo.addrs) == 1 {
opt.SetDirect(true)
}
client = connect(ctx, opt)
})
}
// InitWithOption is a singleton, it will be initialized only once.
func InitWithOption(ctx context.Context, opt *options.ClientOptions) {
once.Do(func() {
tM := reflect.TypeOf(bson.M{})
reg := bson.NewRegistryBuilder().RegisterTypeMapEntry(bsontype.EmbeddedDocument, tM).Build()
opt.SetRegistry(reg)
client = connect(ctx, opt)
})
}
func Close(ctx context.Context) error {
return client.Disconnect(ctx)
}
func Ping(ctx context.Context) error {
return client.Ping(ctx, readpref.Primary())
}
func connect(ctx context.Context, opt *options.ClientOptions) *mongo.Client {
c, err := mongo.Connect(ctx, opt)
if err != nil {
panic(err)
}
return c
}