diff --git a/sys/db/db.go b/sys/db/db.go index 0162aad8a..ff084ef3f 100644 --- a/sys/db/db.go +++ b/sys/db/db.go @@ -16,9 +16,12 @@ type DB struct { } func (this *DB) init() (err error) { - this.mgo, err = mgo.NewSys( + if this.mgo, err = mgo.NewSys( mgo.SetMongodbUrl(this.options.MongodbUrl), mgo.SetMongodbDatabase(this.options.MongodbDatabase), - ) + ); err != nil { + return + } + err = this.checkUserIdInit() return } diff --git a/sys/db/user.go b/sys/db/user.go index 72ded4736..e1cf9c692 100644 --- a/sys/db/user.go +++ b/sys/db/user.go @@ -1,15 +1,52 @@ package db import ( + "context" + "fmt" "go_dreamfactory/pb" + "math/rand" + "time" "github.com/liwei1dao/lego/core" + "github.com/liwei1dao/lego/sys/log" + "go.mongodb.org/mongo-driver/bson" ) const ( //Redis - DB_UserTable core.SqlTable = "user" //会话列表 + DB_UserTable core.SqlTable = "user" //用户表 + DB_UserIdTable core.SqlTable = "userid" //用户id表 ) +//校验数据库初始化工作是否完成 +func (this DB) checkUserIdInit() (err error) { + ctx, _ := context.WithTimeout(context.Background(), time.Second*60) + count, err := this.mgo.CountDocuments(DB_UserIdTable, bson.M{}) + if err != nil || count == 0 { + //批量插入数据 + leng := 1000000 + cIds := make([]interface{}, leng) + for i, _ := range cIds { + cIds[i] = 1000000 + i + } + data := make([]interface{}, leng) + r := rand.New(rand.NewSource(time.Now().Unix())) + n := 0 + for _, i := range r.Perm(leng) { + data[n] = bson.M{"_id": i} + n++ + } + var ( + err error + ) + begin := time.Now() + if _, err = this.mgo.InsertManyByCtx(DB_UserIdTable, ctx, data); err != nil { + return fmt.Errorf("checkUserIdInit err=%s", err.Error()) + } + log.Debugf("checkUserIdInit succ time consuming:%v", time.Now().Sub(begin)) + } + return +} + //更新用户数据到DB func (this *DB) UpdateUser(data *pb.DB_UserData) (err error) {