77 lines
1.8 KiB
Go
77 lines
1.8 KiB
Go
// --------------------------------------------------------------------------------------------------
|
|
// The following code is automatically generated by the mongo-dao-generator tool.
|
|
// Please do not modify this code manually to avoid being overwritten in the next generation.
|
|
// For more tool details, please click the link to view https://github.com/dobyte/mongo-dao-generator
|
|
// --------------------------------------------------------------------------------------------------
|
|
|
|
package internal
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
type Counter struct {
|
|
Columns *CounterColumns
|
|
Database *mongo.Database
|
|
Collection *mongo.Collection
|
|
}
|
|
|
|
type CounterModel struct {
|
|
ID string `bson:"_id"`
|
|
Value int64 `bson:"value"`
|
|
}
|
|
|
|
type CounterColumns struct {
|
|
ID string
|
|
Value string
|
|
}
|
|
|
|
var counterColumns = &CounterColumns{
|
|
ID: "_id",
|
|
Value: "value",
|
|
}
|
|
|
|
func NewCounter(db *mongo.Database) *Counter {
|
|
return &Counter{
|
|
Columns: counterColumns,
|
|
Database: db,
|
|
Collection: db.Collection("counter"),
|
|
}
|
|
}
|
|
|
|
// Incr 自增值
|
|
func (dao *Counter) Incr(ctx context.Context, key string, incr ...int) (int64, error) {
|
|
var (
|
|
upsert = true
|
|
returnDocument = options.After
|
|
counter = &CounterModel{}
|
|
value = 1
|
|
)
|
|
|
|
if len(incr) > 0 {
|
|
if incr[0] == 0 {
|
|
return 0, errors.New("invalid increment value")
|
|
}
|
|
value = incr[0]
|
|
}
|
|
|
|
rst := dao.Collection.FindOneAndUpdate(ctx, bson.M{
|
|
dao.Columns.ID: key,
|
|
}, bson.M{"$inc": bson.M{
|
|
dao.Columns.Value: value,
|
|
}}, &options.FindOneAndUpdateOptions{
|
|
Upsert: &upsert,
|
|
ReturnDocument: &returnDocument,
|
|
})
|
|
|
|
if err := rst.Decode(counter); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
return counter.Value, nil
|
|
}
|