100 lines
2.4 KiB
Go
100 lines
2.4 KiB
Go
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"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
const ( //Redis
|
|
DB_UserTable core.SqlTable = "user" //用户表
|
|
DB_UserIdTable core.SqlTable = "userid" //用户id表
|
|
)
|
|
|
|
type UserId struct {
|
|
UserId uint32 `bson:"_id"`
|
|
}
|
|
|
|
func (this *DB) FindUserByAccount(account string) (*pb.DB_UserData, error) {
|
|
filter := bson.D{
|
|
{"account", account},
|
|
}
|
|
sr := this.mgo.FindOne(DB_UserTable, filter)
|
|
user := &pb.DB_UserData{}
|
|
err := sr.Decode(user)
|
|
return user, err
|
|
}
|
|
|
|
func (this *DB) FindUserById(id uint32) (*pb.DB_UserData, error) {
|
|
filter := bson.D{
|
|
{"_id", id},
|
|
}
|
|
sr := this.mgo.FindOne(DB_UserTable, filter)
|
|
user := &pb.DB_UserData{}
|
|
err := sr.Decode(user)
|
|
return user, err
|
|
}
|
|
|
|
func (this *DB) CreateUser(user *pb.DB_UserData) error {
|
|
userId := &UserId{}
|
|
err := this.mgo.FindOneAndDelete(DB_UserIdTable, bson.M{}).Decode(userId)
|
|
if err != nil {
|
|
log.Errorf("find userId err :%v", err)
|
|
return err
|
|
}
|
|
user.UserId = userId.UserId
|
|
_, err = this.mgo.InsertOne(DB_UserTable, user)
|
|
return err
|
|
}
|
|
|
|
//校验数据库初始化工作是否完成
|
|
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) error {
|
|
err := this.mgo.FindOneAndUpdate(
|
|
DB_UserTable,
|
|
bson.M{"_id": data.UserId},
|
|
bson.M{"$set": bson.M{
|
|
"niceName": data.NiceName,
|
|
"email": data.Email,
|
|
}},
|
|
options.FindOneAndUpdate().SetUpsert(false).SetReturnDocument(options.After),
|
|
).Decode(data)
|
|
return err
|
|
}
|