78 lines
1.6 KiB
Go
78 lines
1.6 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"local/cheqbooq/config"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
type Mongo struct {
|
|
client *mongo.Client
|
|
}
|
|
|
|
func NewMongo() (*Mongo, error) {
|
|
ctx, can := context.WithTimeout(context.Background(), time.Second*5)
|
|
defer can()
|
|
|
|
opt := options.Client()
|
|
opt.ApplyURI(fmt.Sprintf("mongodb://%s", strings.TrimPrefix(config.StoreAddr, "mongodb://")))
|
|
|
|
client, err := mongo.Connect(ctx, opt)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &Mongo{
|
|
client: client,
|
|
}, client.Ping(ctx, nil)
|
|
}
|
|
|
|
func (m *Mongo) Close() error {
|
|
return m.client.Disconnect(context.TODO())
|
|
}
|
|
|
|
func (m *Mongo) Account() *mongo.Collection {
|
|
return m.Database(config.StoreNS).Collection("account")
|
|
}
|
|
|
|
func (m *Mongo) Balance() *mongo.Collection {
|
|
return m.Database(config.StoreNS).Collection("balance")
|
|
}
|
|
|
|
func (m *Mongo) Transaction() *mongo.Collection {
|
|
return m.Database(config.StoreNS).Collection("transaction")
|
|
}
|
|
|
|
func (m *Mongo) Find(c *mongo.Collection, where interface{}, next func() interface{}) error {
|
|
ctx, can := context.WithCancel(context.TODO())
|
|
defer can()
|
|
|
|
cur, err := c.Find(ctx, where, options.Find())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer cur.Close(ctx)
|
|
|
|
for cur.Next(ctx) {
|
|
ptr := next()
|
|
if err := cur.Decode(ptr); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return cur.Err()
|
|
}
|
|
|
|
func (m *Mongo) Upsert(c *mongo.Collection, where, op interface{}) error {
|
|
ctx, can := context.WithCancel(context.TODO())
|
|
defer can()
|
|
|
|
_, err := c.UpdateMany(ctx, where, op, options.Update().SetUpsert(true))
|
|
return err
|
|
}
|