mdbc/index_scope_v2.go

package mdbc

import (
	"context"
	"fmt"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// KeySort 排序方式 类型别名
type KeySort = int

const (
	ASC  KeySort = 1  // 正序排列
	DESC KeySort = -1 // 倒序排列
)

// IndexCard 索引信息
type IndexCard struct {
	Version int32              `bson:"v"`    // 版本信息
	Name    string             `bson:"name"` // 索引名
	Key     map[string]KeySort `bson:"key"`  // 索引方式
}

// IndexList 索引列表
type IndexList []*IndexCard

type IndexScope struct {
	scope      *Scope
	cw         *ctxWrap
	iv         mongo.IndexView
	cursor     *mongo.Cursor
	ims        []mongo.IndexModel
	err        error
	lOpts      *options.ListIndexesOptions
	dOpts      *options.DropIndexesOptions
	cOpts      *options.CreateIndexesOptions
	execResult interface{}
}

func (is *IndexScope) assertErr() {
	if is.err == nil {
		return
	}

	err, ok := is.err.(mongo.CommandError)
	if !ok {
		return
	}
	if err.HasErrorMessage(context.DeadlineExceeded.Error()) {
		is.err = &ErrRequestBroken
	}
}

func (is *IndexScope) doClear() {
	if is.cw != nil && is.cw.cancel != nil {
		is.cw.cancel()
	}
	is.scope.debug = false
	is.scope.execT = 0
}

// doGet 拿到 iv 后续的 增删改查都基于 iv 对象
func (is *IndexScope) doGet() {
	is.iv = db.Collection(is.scope.tableName).Indexes(is.getContext())
}

// doCursor 拿到 cursor
func (is *IndexScope) doCursor() {
	if is.lOpts == nil {
		is.lOpts = &options.ListIndexesOptions{}
	}
	is.cursor, is.err = is.iv.List(is.getContext(), is.lOpts)
	is.assertErr()
}

// SetContext 设定 Context
func (is *IndexScope) SetContext(ctx context.Context) *IndexScope {
	if is.cw == nil {
		is.cw = &ctxWrap{}
	}
	is.cw.ctx = ctx
	return is
}

// SetListIndexesOption 设定 ListIndexesOptions
func (is *IndexScope) SetListIndexesOption(opts options.ListIndexesOptions) *IndexScope {
	is.lOpts = &opts
	return is
}

// SetDropIndexesOption 设定 DropIndexesOptions
func (is *IndexScope) SetDropIndexesOption(opts options.DropIndexesOptions) *IndexScope {
	is.dOpts = &opts
	return is
}

// SetCreateIndexesOption 设定 CreateIndexesOptions
func (is *IndexScope) SetCreateIndexesOption(opts options.CreateIndexesOptions) *IndexScope {
	is.cOpts = &opts
	return is
}

// getContext 获取ctx
func (is *IndexScope) getContext() context.Context {
	return is.cw.ctx
}

// GetIndexList 获取索引列表
func (is *IndexScope) GetIndexList(data *IndexList) error {
	defer is.doClear()
	is.doGet()

	is.doCursor()

	if is.err != nil {
		return is.err
	}

	is.err = is.cursor.All(is.getContext(), data)

	is.assertErr()

	if is.err != nil {
		return is.err
	}

	return nil
}

// GetListCursor 获取结果集的指针
func (is *IndexScope) GetListCursor() (*mongo.Cursor, error) {
	defer is.doClear()

	is.doGet()

	is.doCursor()

	if is.err != nil {
		return nil, is.err
	}

	return is.cursor, nil
}

// DropOne 删除一个索引 name 为索引名称
func (is *IndexScope) DropOne(name string) error {
	defer is.doClear()

	is.doGet()

	is.doDropOne(name)

	if is.err != nil {
		return is.err
	}

	return nil
}

// DropAll 删除所有的索引
func (is *IndexScope) DropAll() error {
	defer is.doClear()

	is.doGet()

	is.doDropAll()
	if is.err != nil {
		return is.err
	}

	return nil
}

func (is *IndexScope) doDropAll() {
	if is.dOpts == nil {
		is.dOpts = &options.DropIndexesOptions{}
	}

	is.execResult, is.err = is.iv.DropAll(is.getContext(), is.dOpts)

	is.assertErr()
}

// doDropOne 删除一个索引
func (is *IndexScope) doDropOne(name string) {
	if is.dOpts == nil {
		is.dOpts = &options.DropIndexesOptions{}
	}

	is.execResult, is.err = is.iv.DropOne(is.getContext(), name, is.dOpts)

	is.assertErr()
}

// CreateOne 创建单个索引
// 在调用本方法前，需要调用 AddIndexModel() 添加
// 当添加了多个 indexModel 后仅会创建第一个索引
// 返回 res 索引名称 err 创建时错误信息
func (is *IndexScope) CreateOne() (res string, err error) {
	defer is.doClear()

	is.doGet()

	is.doCreateOne()

	if is.err != nil {
		return "", is.err
	}

	res, ok := is.execResult.(string)
	if !ok {
		return "", fmt.Errorf("create success but get index name empty")
	}

	return res, nil
}

// CreateMany 创建多个索引，在调用本方法前，需要调用 AddIndexModel() 添加
func (is *IndexScope) CreateMany() (res []string, err error) {
	defer is.doClear()

	is.doGet()

	is.doCreateMany()

	if is.err != nil {
		return nil, is.err
	}

	res, ok := is.execResult.([]string)
	if !ok {
		return nil, fmt.Errorf("create success but get index names empty")
	}

	return res, nil
}

func (is *IndexScope) doCreateOne() {
	if len(is.ims) == 0 {
		is.err = fmt.Errorf("no such index model, you need to use AddIndexModel() to add one")
		return
	}
	if is.cOpts == nil {
		is.cOpts = &options.CreateIndexesOptions{}
	}

	is.execResult, is.err = is.iv.CreateOne(is.getContext(), is.ims[0], is.cOpts)
	is.assertErr()
}

func (is *IndexScope) doCreateMany() {
	if len(is.ims) == 0 {
		is.err = fmt.Errorf("no such index model, you need to use AddIndexModel() to add some")
		return
	}
	if is.cOpts == nil {
		is.cOpts = &options.CreateIndexesOptions{}
	}
	is.execResult, is.err = is.iv.CreateMany(is.getContext(), is.ims, is.cOpts)
	is.assertErr()
}

// AddIndexModels 添加 indexModel
func (is *IndexScope) AddIndexModels(obj ...mongo.IndexModel) *IndexScope {
	if len(obj) == 0 {
		is.err = fmt.Errorf("add index model empty")
		return is
	}
	is.ims = append(is.ims, obj...)
	return is
}