107 lines
3.0 KiB
Go
107 lines
3.0 KiB
Go
|
package mongo
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/zeromicro/go-zero/core/logx"
|
||
|
|
"github.com/zeromicro/go-zero/core/stores/mon"
|
||
|
|
"go.mongodb.org/mongo-driver/v2/bson"
|
||
|
|
"go.mongodb.org/mongo-driver/v2/mongo"
|
||
|
|
"go.mongodb.org/mongo-driver/v2/mongo/options"
|
||
|
|
"go.mongodb.org/mongo-driver/v2/mongo/readpref"
|
||
|
|
)
|
||
|
|
|
||
|
|
type DocumentDB struct {
|
||
|
|
Mon *mon.Model
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewDocumentDB(config *Conf, collection string, opts ...mon.Option) (DocumentDBUseCase, error) {
|
||
|
|
if config == nil {
|
||
|
|
return nil, fmt.Errorf("mongo: config is nil")
|
||
|
|
}
|
||
|
|
if collection == "" {
|
||
|
|
return nil, fmt.Errorf("mongo: collection is required")
|
||
|
|
}
|
||
|
|
|
||
|
|
connectionURI, err := buildConnectionURI(*config)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
opts = append(opts, InitMongoOptions(*config))
|
||
|
|
logx.Infof("[DocumentDB] connecting to %s db=%s coll=%s", redactConnectionURI(connectionURI), config.Database, collection)
|
||
|
|
|
||
|
|
client, err := mon.NewModel(connectionURI, config.Database, collection, opts...)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("mongo: new model: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
timeout := time.Duration(config.ConnectTimeoutMs) * time.Millisecond
|
||
|
|
if timeout <= 0 {
|
||
|
|
timeout = 10 * time.Second
|
||
|
|
}
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
||
|
|
defer cancel()
|
||
|
|
|
||
|
|
if err := client.Collection.Database().Client().Ping(ctx, readpref.Primary()); err != nil {
|
||
|
|
return nil, fmt.Errorf("mongo: ping primary: %w", err)
|
||
|
|
}
|
||
|
|
logx.Infof("[DocumentDB] connected")
|
||
|
|
|
||
|
|
return &DocumentDB{Mon: client}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (document *DocumentDB) PopulateIndex(ctx context.Context, key string, sort int32, unique bool) error {
|
||
|
|
return document.createIndex(ctx, []string{key}, []int32{sort}, unique, nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (document *DocumentDB) PopulateTTLIndex(ctx context.Context, key string, sort int32, unique bool, ttl int32) error {
|
||
|
|
return document.createIndex(ctx, []string{key}, []int32{sort}, unique, options.Index().SetExpireAfterSeconds(ttl))
|
||
|
|
}
|
||
|
|
|
||
|
|
func (document *DocumentDB) PopulateMultiIndex(ctx context.Context, keys []string, sorts []int32, unique bool) error {
|
||
|
|
if len(keys) != len(sorts) {
|
||
|
|
return fmt.Errorf("mongo: keys and sorts length mismatch")
|
||
|
|
}
|
||
|
|
return document.createIndex(ctx, keys, sorts, unique, nil)
|
||
|
|
}
|
||
|
|
|
||
|
|
func (document *DocumentDB) createIndex(
|
||
|
|
ctx context.Context,
|
||
|
|
keys []string,
|
||
|
|
sorts []int32,
|
||
|
|
unique bool,
|
||
|
|
indexOpt *options.IndexOptionsBuilder,
|
||
|
|
) error {
|
||
|
|
c := document.Mon.Collection
|
||
|
|
index := yieldIndexModel(keys, sorts, unique, indexOpt)
|
||
|
|
ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
_, err := c.Indexes().CreateOne(ctx, index)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("mongo: create index: %w", err)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (document *DocumentDB) GetClient() *mon.Model {
|
||
|
|
return document.Mon
|
||
|
|
}
|
||
|
|
|
||
|
|
func yieldIndexModel(keys []string, sorts []int32, unique bool, indexOpt *options.IndexOptionsBuilder) mongo.IndexModel {
|
||
|
|
setKeysDoc := bson.D{}
|
||
|
|
for i, key := range keys {
|
||
|
|
setKeysDoc = append(setKeysDoc, bson.E{Key: key, Value: sorts[i]})
|
||
|
|
}
|
||
|
|
if indexOpt == nil {
|
||
|
|
indexOpt = options.Index()
|
||
|
|
}
|
||
|
|
indexOpt.SetUnique(unique)
|
||
|
|
return mongo.IndexModel{
|
||
|
|
Keys: setKeysDoc,
|
||
|
|
Options: indexOpt,
|
||
|
|
}
|
||
|
|
}
|