s_cherry_e/sr/internal/component/db/component.go

package db

import (
	"context"
	dbFacade "f1-game/internal/component/db/facade"
	dbMigrator "f1-game/internal/component/db/migrator"
	dbQueue "f1-game/internal/component/db/queue"
	"fmt"

	cstring "github.com/cherry-game/cherry/extend/string"
	clog "github.com/cherry-game/cherry/logger"
	cmongo "github.com/cherry-game/components/mongo"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type Component struct {
	*cmongo.Component
	name            string
	db              *mongo.Database
	tables          []dbFacade.Table
	migratorScripts []dbMigrator.IScript
	dbQueue         *dbQueue.TableQueue
	onLoadFuncs     []func()
}

func (p *Component) Name() string {
	return p.name
}

func New(name string) *Component {
	return &Component{
		name:      name,
		Component: cmongo.NewComponent(),
	}
}

func (p *Component) Init() {
	p.Component.Init()

	p.db = p.loadMongo("mongo_id")
	if p.db == nil {
		panic("db id not found.")
	}

	p.checkIndexes()
	p.initMigrator()
	p.initDBQueue()

	clog.Infof("[db_id = %s] database is init.", p.db.Name())

	for _, fn := range p.onLoadFuncs {
		if fn != nil {
			fn()
		}
	}
}

// 添加映射表
func (p *Component) AddTable(table ...dbFacade.Table) {
	p.tables = append(p.tables, table...)
}

// 添加DB脚本
func (p *Component) AddScript(script ...dbMigrator.IScript) {
	p.migratorScripts = append(p.migratorScripts, script...)
}

// 初始化db脚本
func (p *Component) initMigrator() {
	if len(p.migratorScripts) > 0 {
		migrator := &dbMigrator.Migrator{}
		migrator.Registers(p.migratorScripts)
		migrator.Run(p.App().NodeID(), p.db, p.tables)
	}
}

// 初始化db队列
func (p *Component) initDBQueue() {
	queueConfig := dbQueue.QueueConfig{}
	if err := queueConfig.Load(); err != nil {
		clog.Panic(err)
	}

	p.dbQueue = dbQueue.NewTableQueue(p.db)

	for _, table := range p.tables {
		if !table.EnableQueue() {
			continue
		}

		opts := queueConfig.BuildOpts(table.TableName())
		p.dbQueue.AddTable(table, opts...)
	}

	p.dbQueue.Start()
}

// 检查每个Collection的索引
func (p *Component) checkIndexes() {
	// 构建 获取索引的函数
	findIndexFunc := func(collection *mongo.Collection) map[string]bool {
		indexMaps := map[string]bool{}
		// 读取collection的索引列表
		cursor, err := collection.Indexes().List(context.TODO())
		if err != nil {
			clog.Fatal(err)
		}

		defer cursor.Close(context.TODO())

		for cursor.Next(context.TODO()) {
			var index bson.M
			if err := cursor.Decode(&index); err != nil {
				clog.Fatal(err)
			}

			indexKey := cstring.ToString(index["name"])
			if indexKey != "_id_" {
				indexMaps[indexKey] = true
			}
		}

		return indexMaps
	}

	// 构建 创建索引的函数包装
	createIndexFunc := func(collection *mongo.Collection, indexObject mongo.IndexModel) {
		indexName, err := collection.Indexes().CreateOne(context.TODO(), indexObject)
		if err != nil {
			clog.Fatalf("Create index error. [table = %s, err = %v]", collection.Name(), err)
		} else {
			clog.Infof("Create index. [table = %s, indexName = %s]", collection.Name(), indexName)
		}
	}

	buildIndexNameFunc := func(d bson.D) string {
		var (
			key  string
			size = len(d)
		)

		for i, v := range d {
			key += fmt.Sprintf("%s_%v", v.Key, v.Value)
			if i+1 < size {
				key += "_"
			}
		}
		return key
	}

	for _, table := range p.tables {
		tableIndexes := table.Indexes()
		if len(tableIndexes) < 1 {
			continue
		}

		collection := p.Collection(table.TableName())

		// 获取表已创建的索引
		indexMaps := findIndexFunc(collection)
		if len(indexMaps) < 1 {
			// 全量创建已定义的索引
			for _, indexObject := range tableIndexes {
				createIndexFunc(collection, indexObject)
			}
		} else {
			// 创建新增的索引
			for _, indexObject := range tableIndexes {
				bsonD := indexObject.Keys.(bson.D)
				key := buildIndexNameFunc(bsonD)
				if !indexMaps[key] {
					createIndexFunc(collection, indexObject)
				}
			}
		}
	}
}

func (p *Component) DBQueue() *dbQueue.TableQueue {
	return p.dbQueue
}

// 设置加载函数(组件初始化后执行)
func (p *Component) SetOnLoadFunc(fn ...func()) {
	p.onLoadFuncs = append(p.onLoadFuncs, fn...)
}

func (p *Component) OnStop() {
	p.dbQueue.Stop()
}

// 根据节点配置获取mongo db实
func (p *Component) loadMongo(id string) *mongo.Database {
	dbID := p.App().Settings().GetConfig("mongo_id_list").GetString(id)
	return p.Component.GetDb(dbID)
}

// 获取db对象
func (p *Component) DB() *mongo.Database {
	return p.db
}

// 获取集合对象
func (p *Component) Collection(name string, opts ...options.Lister[options.CollectionOptions]) *mongo.Collection {
	return p.db.Collection(name, opts...)
}

// 查找一个对象
func (p *Component) FindOne(retTable dbFacade.Table, filter bson.M, opts ...options.Lister[options.FindOneOptions]) error {
	if len(opts) < 1 {
		opts = append(opts, options.FindOne())
	}

	return p.Collection(retTable.TableName()).FindOne(context.TODO(), filter, opts...).Decode(retTable)
}

// 分页
func Pagination[T dbFacade.Table](collection *mongo.Collection, ctx context.Context, filter bson.M, projection bson.M, sortOpts bson.D,
	pageNum, pageSize int32, isTotal bool) ([]T, int64, error) {
	if pageNum < 1 {
		pageNum = 1
	}

	if pageSize < 1 {
		pageSize = 1
	}

	if len(sortOpts) < 1 {
		sortOpts = bson.D{{
			Key:   "_id",
			Value: 1,
		}}
	}

	if filter == nil {
		filter = bson.M{}
	}

	if projection == nil {
		projection = bson.M{}
	}

	var (
		skipCount = (pageNum - 1) * pageSize
		findOpts  = options.Find().SetSkip(int64(skipCount)).SetLimit(int64(pageSize)).SetSort(sortOpts).SetProjection(projection)
	)

	cursor, err := collection.Find(ctx, filter, findOpts)
	if err != nil {
		return nil, 0, err
	}

	defer cursor.Close(ctx)

	var results []T

	for cursor.Next(ctx) {
		var item T
		err := cursor.Decode(&item)
		if err != nil {
			return nil, 0, err
		}
		results = append(results, item)
	}

	if err := cursor.Err(); err != nil {
		return nil, 0, err
	}

	var total int64 = 0

	if isTotal {
		total, err = collection.CountDocuments(ctx, filter)
		if err != nil {
			return nil, 0, err
		}
	}

	return results, total, nil
}