mdbc/count.go

72 lines
1.4 KiB
Go

package mdbc
import (
"context"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo/options"
"reflect"
)
type countScope struct {
scope *scope
mdbc *mdbc
ctx context.Context
limit int64
skip int64
filter interface{}
opts options.CountOptions
}
// SetContext set operate context, default timeout 30s
func (cs *countScope) SetContext(ctx context.Context) *countScope {
cs.ctx = ctx
return cs
}
// SetLimit set filter limit size
func (cs *countScope) SetLimit(limit int64) *countScope {
cs.limit = limit
return cs
}
// SetSkip set filter skip size
func (cs *countScope) SetSkip(skip int64) *countScope {
cs.skip = skip
return cs
}
func (cs *countScope) SetFilter(filter interface{}) *countScope {
if filter == nil {
cs.filter = bson.M{}
return cs
}
v := reflect.ValueOf(filter)
if v.Kind() == reflect.Ptr || v.Kind() == reflect.Map || v.Kind() == reflect.Slice {
if v.IsNil() {
cs.filter = bson.M{}
}
}
cs.filter = filter
return cs
}
func (cs *countScope) mergeOptions() {
if cs.ctx == nil {
cs.ctx = context.Background()
}
if cs.filter == nil {
cs.filter = bson.M{}
}
if cs.skip != 0 {
cs.opts.Skip = &cs.skip
}
if cs.limit != 0 {
cs.opts.Limit = &cs.limit
}
}
func (cs *countScope) Count() (int64, error) {
cs.mergeOptions()
return cs.mdbc.database.Collection(cs.scope.tableName).CountDocuments(cs.ctx, cs.filter, &cs.opts)
}