393 lines
12 KiB
Go
393 lines
12 KiB
Go
package db
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"reflect"
|
|
"time"
|
|
|
|
"91porn-server/common/log"
|
|
"91porn-server/common/stderr"
|
|
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/bson/primitive"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
var (
|
|
enableSortMCheck = true
|
|
bsonDType = reflect.TypeOf(bson.D{})
|
|
bsonDPtrType = reflect.TypeOf(&bson.D{})
|
|
bsonEType = reflect.TypeOf(bson.E{})
|
|
bsonEPtrType = reflect.TypeOf(&bson.E{})
|
|
bsonMType = reflect.TypeOf(bson.M{})
|
|
bsonMPtrType = reflect.TypeOf(&bson.M{})
|
|
)
|
|
|
|
var skipErrors = []error{mongo.ErrNoDocuments}
|
|
|
|
func handleDbError(err error) error {
|
|
for _, e := range skipErrors {
|
|
if err == e {
|
|
return nil
|
|
}
|
|
}
|
|
if err != nil {
|
|
if mongo.IsDuplicateKeyError(err) {
|
|
return stderr.InsertExistError
|
|
}
|
|
}
|
|
return err
|
|
}
|
|
|
|
func sortCheck(sort interface{}) error {
|
|
if !enableSortMCheck {
|
|
return nil
|
|
}
|
|
if sort == nil {
|
|
return nil
|
|
}
|
|
typ := reflect.TypeOf(sort)
|
|
var m bson.M
|
|
switch typ {
|
|
case bsonMType:
|
|
//log.Warn("mongo sort use bson.M use bson.D instead", log.Any("sort", sort))
|
|
m, _ = sort.(bson.M)
|
|
case bsonMPtrType:
|
|
//log.Warn("mongo sort use *bson.M use bson.D instead", log.Any("sort", sort))
|
|
pm, _ := sort.(*bson.M)
|
|
m = *pm
|
|
case bsonDType, bsonDPtrType, bsonEType, bsonEPtrType:
|
|
return nil
|
|
default:
|
|
log.Warn("sort use unknown sort type please check", log.Any("sort", sort), log.Any("typ", typ))
|
|
return errors.New("mongo error sort type")
|
|
}
|
|
if len(m) > 1 {
|
|
return errors.New("mongo error sort, use bson.M and len(sort) > 1")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// CreateIndex 创建数据索引.options 在index里面创建
|
|
func (u *MongoTool) CreateIndex(models []mongo.IndexModel) ([]string, error) {
|
|
//return nil, nil
|
|
return u.coll.Indexes().CreateMany(u.ctx, models)
|
|
}
|
|
|
|
// DropIndex 删除数据索引.options 在index里面创建
|
|
func (u *MongoTool) DropIndex(indexname string) error {
|
|
if _, err := u.coll.Indexes().DropOne(u.ctx, indexname); err != nil {
|
|
log.Error(fmt.Sprintf("drop indexes error %+v", err))
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// DropIndexIfExists 删除指定索引；索引或所在集合尚未创建时均视为成功。
|
|
func (u *MongoTool) DropIndexIfExists(indexname string) error {
|
|
if _, err := u.coll.Indexes().DropOne(u.ctx, indexname); err != nil {
|
|
var commandErr mongo.CommandError
|
|
// 27=IndexNotFound(索引不存在)、26=NamespaceNotFound(集合/库尚未创建)：目标索引本就不存在，视为成功。
|
|
if errors.As(err, &commandErr) && (commandErr.Code == 27 || commandErr.Code == 26) {
|
|
return nil
|
|
}
|
|
log.Error(fmt.Sprintf("drop indexes error %+v", err))
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// InsertOne 插入单条信息
|
|
func (u *MongoTool) InsertOne(document interface{}, opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
|
|
insertResult, err := u.coll.InsertOne(u.ctx, document, opts...)
|
|
return insertResult, handleDbError(err)
|
|
}
|
|
|
|
// InsertMany 批量插入信息
|
|
func (u *MongoTool) InsertMany(documents interface{}, opts ...*options.InsertManyOptions) (*mongo.InsertManyResult, error) {
|
|
if err := validInterfaceSlice(documents); err != nil {
|
|
return nil, err
|
|
}
|
|
res := setTimeForSlice(documents)
|
|
return u.coll.InsertMany(u.ctx, res, opts...)
|
|
}
|
|
|
|
// Find 查询多条数据
|
|
func (u *MongoTool) Find(model interface{}, filter bson.M, opts ...*options.FindOptions) error {
|
|
if err := validInterfaceSlice(model); err != nil {
|
|
return err
|
|
}
|
|
cur, err := u.FindCursor(filter, opts...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return handleDbError(cur.All(u.ctx, model))
|
|
}
|
|
|
|
// FindCursor 查询多条数据并返回游标。
|
|
// Cursor 不是并发安全的，调用方必须在完成或失败后关闭它。
|
|
func (u *MongoTool) FindCursor(filter bson.M, opts ...*options.FindOptions) (*mongo.Cursor, error) {
|
|
for _, opt := range opts {
|
|
if opt == nil {
|
|
continue
|
|
}
|
|
if opt.Limit != nil && *opt.Limit > 1000 {
|
|
fmt.Println("limit beyond 1000 ==================>", *opt.Limit)
|
|
}
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
}
|
|
return u.coll.Find(u.ctx, filter, opts...)
|
|
}
|
|
|
|
// FindOne 单条查询
|
|
func (u *MongoTool) FindOne(model interface{}, filter bson.M, opts ...*options.FindOneOptions) error {
|
|
for _, opt := range opts {
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return handleDbError(u.coll.FindOne(u.ctx, filter, opts...).Decode(model))
|
|
}
|
|
|
|
// FindOneAndDelete 单条查询并删除
|
|
func (u *MongoTool) FindOneAndDelete(model interface{}, filter bson.M, opts ...*options.FindOneAndDeleteOptions) error {
|
|
for _, opt := range opts {
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return handleDbError(u.coll.FindOneAndDelete(u.ctx, filter, opts...).Decode(model))
|
|
}
|
|
|
|
// FindOneAndReplace 单条查询 rd set to Before 表示返回原始数据, set to After 表示返回替换后的数据
|
|
func (u *MongoTool) FindOneAndReplace(model interface{}, filter bson.M, replacement bson.M, opts ...*options.FindOneAndReplaceOptions) error {
|
|
for _, opt := range opts {
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return handleDbError(u.coll.FindOneAndReplace(u.ctx, filter, replacement, opts...).Decode(model))
|
|
}
|
|
|
|
// FindOneAndUpdate 单条查询 rd set to Before 表示返回原始数据, set to After 表示返回更新后的数据 默认为返回更新后的数据
|
|
func (u *MongoTool) FindOneAndUpdate(model interface{}, filter bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error {
|
|
for _, opt := range opts {
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
setReturn := false
|
|
for _, opt := range opts {
|
|
if opt != nil && opt.ReturnDocument != nil {
|
|
setReturn = true
|
|
break
|
|
}
|
|
}
|
|
if !setReturn {
|
|
after := options.After
|
|
opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &after})
|
|
}
|
|
return handleDbError(u.coll.FindOneAndUpdate(u.ctx, filter, update, opts...).Decode(model))
|
|
}
|
|
|
|
// FindOneAndUpsert 单条查询 匹配到数据更新,未匹配到数据则upsert
|
|
func (u *MongoTool) FindOneAndUpsert(model interface{}, filter bson.M, update bson.M, opts ...*options.FindOneAndUpdateOptions) error {
|
|
upsert := true
|
|
var beforeOrAfter options.ReturnDocument
|
|
for _, opt := range opts {
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
setReturn := false
|
|
for _, opt := range opts {
|
|
if opt != nil && opt.ReturnDocument != nil {
|
|
beforeOrAfter = *opt.ReturnDocument
|
|
setReturn = true
|
|
break
|
|
}
|
|
}
|
|
if !setReturn {
|
|
beforeOrAfter = options.After
|
|
}
|
|
opts = append(opts, &options.FindOneAndUpdateOptions{ReturnDocument: &beforeOrAfter, Upsert: &upsert})
|
|
return handleDbError(u.coll.FindOneAndUpdate(u.ctx, filter, update, opts...).Decode(model))
|
|
}
|
|
|
|
func (u *MongoTool) FindOneAndUpdateReturnTiny(bind interface{}, query bson.M, update bson.M, afterDoc bool, opts ...*options.FindOneAndUpdateOptions) error {
|
|
if afterDoc {
|
|
opts = append(opts, options.FindOneAndUpdate().SetReturnDocument(options.After))
|
|
} else {
|
|
opts = append(opts, options.FindOneAndUpdate().SetReturnDocument(options.Before))
|
|
}
|
|
result := handleDbError(u.coll.FindOneAndUpdate(u.ctx, query, update, opts...).Decode(bind))
|
|
return result
|
|
}
|
|
|
|
// FindOneByID 通过id查找一条数据
|
|
func (u *MongoTool) FindOneByID(model interface{}, id primitive.ObjectID, opts ...*options.FindOneOptions) error {
|
|
for _, opt := range opts {
|
|
if opt.Sort != nil {
|
|
if err := sortCheck(opt.Sort); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return handleDbError(u.coll.FindOne(u.ctx, bson.M{"_id": id}, opts...).Decode(model))
|
|
}
|
|
|
|
// Aggregate 聚合查找数据
|
|
func (u *MongoTool) Aggregate(model interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error {
|
|
if err := validInterfaceSlice(model); err != nil {
|
|
return err
|
|
}
|
|
cur, err := u.coll.Aggregate(u.ctx, pipeline, opts...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return handleDbError(cur.All(u.ctx, model))
|
|
}
|
|
|
|
// AggregateDecode 聚合.Decode
|
|
func (u *MongoTool) AggregateDecode(model interface{}, pipeline []bson.M, opts ...*options.AggregateOptions) error {
|
|
cur, err := u.coll.Aggregate(u.ctx, pipeline, opts...)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if cur.Next(u.ctx) {
|
|
return handleDbError(cur.Decode(model))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Distinct 去重查询
|
|
func (u *MongoTool) Distinct(fieldName string, filter bson.M, opts ...*options.DistinctOptions) ([]interface{}, error) {
|
|
return u.coll.Distinct(u.ctx, fieldName, filter, opts...)
|
|
}
|
|
|
|
// DeleteOne 删除一条数据
|
|
func (u *MongoTool) DeleteOne(filter bson.M, opt ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
|
|
return u.coll.DeleteOne(u.ctx, filter, opt...)
|
|
}
|
|
|
|
// DeleteMany 删除多条数据
|
|
func (u *MongoTool) DeleteMany(filter bson.M, opt ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
|
|
return u.coll.DeleteMany(u.ctx, filter, opt...)
|
|
}
|
|
|
|
// DeleteById 根据ID删除数据单条数据
|
|
func (u *MongoTool) DeleteById(id primitive.ObjectID) (*mongo.DeleteResult, error) {
|
|
return u.coll.DeleteOne(u.ctx, bson.M{"_id": id})
|
|
}
|
|
|
|
// UpdateOne 更新单条数据
|
|
func (u *MongoTool) UpdateOne(filter bson.M, update interface{}) (*mongo.UpdateResult, error) {
|
|
return u.coll.UpdateOne(u.ctx, filter, update)
|
|
}
|
|
|
|
// UpdateMany 修改多条数据
|
|
func (u *MongoTool) UpdateMany(filter bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
|
return u.coll.UpdateMany(u.ctx, filter, update)
|
|
}
|
|
|
|
// UpsertMany 或者修改或者插入多条数据
|
|
func (u *MongoTool) UpsertMany(filter bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
|
return u.coll.UpdateMany(u.ctx, filter, update, options.Update().SetUpsert(true))
|
|
}
|
|
|
|
// UpsertOne 或者修改或者插入一条数据
|
|
func (u *MongoTool) UpsertOne(filter bson.M, update bson.M) (*mongo.UpdateResult, error) {
|
|
return u.coll.UpdateOne(u.ctx, filter, update, options.Update().SetUpsert(true))
|
|
}
|
|
|
|
// UpdateOneForSet 修改一条数据 【根据修改数据中集合类型字段】
|
|
func (u *MongoTool) UpdateOneForSet(filter bson.M, update bson.D) (*mongo.UpdateResult, error) {
|
|
return u.coll.UpdateOne(u.ctx, filter, update)
|
|
}
|
|
|
|
// Count 获取数量
|
|
func (u *MongoTool) Count(filter interface{}, opts ...*options.CountOptions) (int64, error) {
|
|
if reflect.TypeOf(filter).Kind() == reflect.Slice {
|
|
cur, err := u.coll.Aggregate(u.ctx, filter)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
var count int64 = 0
|
|
for cur.Next(context.TODO()) {
|
|
count++
|
|
}
|
|
return count, nil
|
|
}
|
|
return u.coll.CountDocuments(u.ctx, filter, opts...)
|
|
}
|
|
|
|
// EstimateCount 获取数量
|
|
func (u *MongoTool) EstimateCount(opts ...*options.EstimatedDocumentCountOptions) (int64, error) {
|
|
return u.coll.EstimatedDocumentCount(u.ctx, opts...)
|
|
}
|
|
|
|
// Bulk Bulk
|
|
func (u *MongoTool) Bulk(models []mongo.WriteModel, opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) {
|
|
return u.coll.BulkWrite(u.ctx, models, opts...)
|
|
}
|
|
|
|
// Exists 是否存在数据
|
|
func (u *MongoTool) Exists(filter interface{}, opts ...*options.FindOneOptions) (bool, error) {
|
|
var limit int64 = 1
|
|
lo := &options.CountOptions{Limit: &limit}
|
|
n, err := u.Count(filter, lo)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
return n > 0, nil
|
|
}
|
|
|
|
func setTimeForSlice(docs interface{}) []interface{} {
|
|
arr := reflect.ValueOf(docs)
|
|
if arr.Kind() == reflect.Ptr {
|
|
arr = reflect.ValueOf(docs).Elem()
|
|
}
|
|
result := make([]interface{}, arr.Len())
|
|
for i := 0; i < arr.Len(); i++ {
|
|
ele := arr.Index(i)
|
|
now := time.Now()
|
|
if ma := ele.FieldByName("UpdatedAt"); ma.IsValid() {
|
|
ma.Set(reflect.ValueOf(now))
|
|
}
|
|
if ca := ele.FieldByName("CreatedAt"); ca.IsValid() {
|
|
ca.Set(reflect.ValueOf(now))
|
|
}
|
|
result[i] = ele.Interface()
|
|
}
|
|
return result
|
|
}
|
|
|
|
func validInterfaceSlice(bind interface{}) error {
|
|
t := reflect.TypeOf(bind)
|
|
k := t.Kind()
|
|
if t.Kind() == reflect.Ptr {
|
|
k = t.Elem().Kind()
|
|
}
|
|
if k != reflect.Slice {
|
|
return stderr.MustSliceOrSlicePtr
|
|
}
|
|
return nil
|
|
}
|