112 lines
2.7 KiB
Go
112 lines
2.7 KiB
Go
package mdbc
|
|
|
|
import (
|
|
"context"
|
|
"github.com/sirupsen/logrus"
|
|
"gitter.top/coco/lormatter"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
"go.mongodb.org/mongo-driver/mongo/readpref"
|
|
"google.golang.org/protobuf/proto"
|
|
"reflect"
|
|
"time"
|
|
)
|
|
|
|
func init() {
|
|
formatter := &lormatter.Formatter{}
|
|
logrus.SetFormatter(formatter)
|
|
logrus.SetReportCaller(true)
|
|
}
|
|
|
|
type Config struct {
|
|
uri string
|
|
Username string `json:"username" yaml:"username"`
|
|
Password string `json:"password" yaml:"password"`
|
|
Address string `json:"address" yaml:"address"`
|
|
Port int `json:"port" yaml:"port"`
|
|
ProcessTimeout int `json:"process_timeout" yaml:"process_timeout"`
|
|
DbName string `json:"db_name" yaml:"db_name"`
|
|
ReadPref ReadPref `json:"read_pref" yaml:"read_pref"`
|
|
}
|
|
|
|
func (c *Config) merge() {
|
|
if c.Address == "" {
|
|
c.Address = "localhost"
|
|
}
|
|
if c.Port == 0 {
|
|
c.Port = 27017
|
|
}
|
|
if c.ProcessTimeout == 0 {
|
|
c.ProcessTimeout = 30
|
|
}
|
|
}
|
|
|
|
type mdbc struct {
|
|
client *mongo.Client
|
|
database *mongo.Database
|
|
readpref *readpref.ReadPref
|
|
}
|
|
|
|
func NewMDBC(config *Config) *mdbc {
|
|
var driver = new(mdbc)
|
|
var clientOpts = options.Client()
|
|
config.merge()
|
|
clientOpts.ApplyURI(applyURI(config.Address, config.Port))
|
|
if config.Username != "" && config.Password != "" {
|
|
credential := options.Credential{
|
|
Username: config.Username,
|
|
Password: config.Password,
|
|
}
|
|
clientOpts.SetAuth(credential)
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(config.ProcessTimeout)*time.Second)
|
|
defer cancel()
|
|
client, err := mongo.Connect(ctx, clientOpts)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
driver.readpref = readPref(config.ReadPref)
|
|
if err := client.Ping(context.Background(), driver.readpref); err != nil {
|
|
panic(err)
|
|
}
|
|
driver.client = client
|
|
if config.DbName != "" {
|
|
driver.database = client.Database(config.DbName)
|
|
}
|
|
return driver
|
|
}
|
|
|
|
// GetClient get raw mongo client
|
|
func (m *mdbc) GetClient() *mongo.Client {
|
|
return m.client
|
|
}
|
|
|
|
// GetDatabase get raw mongo database
|
|
func (m *mdbc) GetDatabase() *mongo.Database {
|
|
return m.database
|
|
}
|
|
|
|
// BindModel mapping protobuf Message to mongo Collection
|
|
func (m *mdbc) BindModel(prototype any) Collection {
|
|
if prototype == nil {
|
|
panic("model can not be nil")
|
|
}
|
|
if _, impled := prototype.(tabler); !impled {
|
|
panic("model does not implement tabler interface")
|
|
}
|
|
protomsg, impled := prototype.(proto.Message)
|
|
if !impled {
|
|
panic("model does not implement message interface")
|
|
}
|
|
scope := &scope{
|
|
&model{
|
|
Type: protomsg,
|
|
modelKind: reflect.TypeOf(protomsg),
|
|
modelName: string(protomsg.ProtoReflect().Descriptor().FullName()),
|
|
tableName: prototype.(tabler).TableName(),
|
|
},
|
|
m,
|
|
}
|
|
return scope
|
|
}
|