48 lines
1.2 KiB
Go
Executable File
48 lines
1.2 KiB
Go
Executable File
package mongo
|
|
|
|
import (
|
|
"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 Conf) mon.Option {
|
|
return func(opts *options.ClientOptions) {
|
|
opts.SetMaxPoolSize(cfg.MaxPoolSize)
|
|
opts.SetMinPoolSize(cfg.MinPoolSize)
|
|
opts.SetMaxConnIdleTime(cfg.MaxConnIdleTime)
|
|
opts.SetCompressors([]string{"snappy"})
|
|
}
|
|
}
|