52 lines
1.4 KiB
Go
52 lines
1.4 KiB
Go
|
package mongo
|
||
|
|
|
||
|
|
import (
|
||
|
|
"app-cloudep-trade-service/internal/config"
|
||
|
|
"reflect"
|
||
|
|
|
||
|
|
"github.com/shopspring/decimal"
|
||
|
|
"github.com/zeromicro/go-zero/core/stores/mon"
|
||
|
|
"go.mongodb.org/mongo-driver/bson"
|
||
|
|
"go.mongodb.org/mongo-driver/bson/bsoncodec"
|
||
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
||
|
|
)
|
||
|
|
|
||
|
|
type TypeCodec struct {
|
||
|
|
ValueType reflect.Type
|
||
|
|
Encoder bsoncodec.ValueEncoder
|
||
|
|
Decoder bsoncodec.ValueDecoder
|
||
|
|
}
|
||
|
|
|
||
|
|
// WithTypeCodec registers TypeCodecs to convert custom types.
|
||
|
|
func WithTypeCodec(typeCodecs ...TypeCodec) mon.Option {
|
||
|
|
return func(c *options.ClientOptions) {
|
||
|
|
registry := bson.NewRegistry()
|
||
|
|
for _, v := range typeCodecs {
|
||
|
|
registry.RegisterTypeEncoder(v.ValueType, v.Encoder)
|
||
|
|
registry.RegisterTypeDecoder(v.ValueType, v.Decoder)
|
||
|
|
}
|
||
|
|
c.SetRegistry(registry)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// SetCustomDecimalType force convert primitive.Decimal128 to decimal.Decimal.
|
||
|
|
func SetCustomDecimalType() mon.Option {
|
||
|
|
return WithTypeCodec(TypeCodec{
|
||
|
|
ValueType: reflect.TypeOf(decimal.Decimal{}),
|
||
|
|
Encoder: &MgoDecimal{},
|
||
|
|
Decoder: &MgoDecimal{},
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func InitMongoOptions(cfg config.Config) mon.Option {
|
||
|
|
return func(opts *options.ClientOptions) {
|
||
|
|
opts.SetMaxPoolSize(cfg.Mongo.MaxPoolSize)
|
||
|
|
opts.SetMinPoolSize(cfg.Mongo.MinPoolSize)
|
||
|
|
opts.SetMaxConnIdleTime(cfg.Mongo.MaxConnIdleTime)
|
||
|
|
opts.SetCompressors([]string{"snappy"})
|
||
|
|
// opts.SetReplicaSet(cfg.Mongo.ReplicaName)
|
||
|
|
// opts.SetWriteConcern(writeconcern.W1())
|
||
|
|
// opts.SetReadPreference(readpref.Primary())
|
||
|
|
}
|
||
|
|
}
|