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 }