Skip to content

Commit fd11da7

Browse files
xyy0411xyy0411
committed
将 ensureTable 封装到钩子中，减少冗余代码
1 parent 87bd28f commit fd11da7

File tree

2 files changed

+40
-60
lines changed

2 files changed

+40
-60
lines changed

‎niu/db.go‎

Lines changed: 37 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,22 @@ import (
44
"fmt"
55
"github.com/RomiChan/syncx"
66
"github.com/jinzhu/gorm"
7+
"sync"
78
)
89

910
var (
1011
migratedGroups = syncx.Map[string, bool]{} // key: string, value: bool
12+
tableHooks []tableHook
13+
hooksMtx sync.RWMutex
1114
)
1215

13-
func ensureUserInfoTable[T userInfo | AuctionInfo](gid int64, prefix string) error {
16+
type tableHook func(gid int64) error
17+
18+
type model struct {
19+
*gorm.DB
20+
}
21+
22+
func ensureTable[T userInfo | AuctionInfo](gid int64, prefix string) error {
1423
table := fmt.Sprintf("group_%d_%s_info", gid, prefix)
1524
if _, ok := migratedGroups.Load(table); ok {
1625
return nil
@@ -25,9 +34,34 @@ func ensureUserInfoTable[T userInfo | AuctionInfo](gid int64, prefix string) err
2534
return nil
2635
}
2736

37+
func ensureUserInfo(gid int64) error {
38+
return ensureTable[userInfo](gid, ur)
39+
}
40+
41+
func ensureAuctionInfo(gid int64) error {
42+
return ensureTable[AuctionInfo](gid, ac)
43+
}
44+
45+
// registerTableHook 注册钩子
46+
func registerTableHook(h ...tableHook) {
47+
hooksMtx.Lock()
48+
defer hooksMtx.Unlock()
49+
tableHooks = append(tableHooks, h...)
50+
}
51+
2852
// TableFor 大写是为了防止数据操作哪里有问题留个保底可以在zbp的项目里直接改
29-
func TableFor(gid int64, prefix string) *gorm.DB {
30-
return db.Table(fmt.Sprintf("group_%d_%s_info", gid, prefix))
53+
func TableFor(gid int64, prefix string) *model {
54+
// 先执行钩子
55+
hooksMtx.RLock()
56+
for _, h := range tableHooks {
57+
if err := h(gid); err != nil {
58+
panic(fmt.Sprintf("执行表钩子失败: %v", err))
59+
}
60+
}
61+
hooksMtx.RUnlock()
62+
63+
tableName := fmt.Sprintf("group_%d_%s_info", gid, prefix)
64+
return &model{db.Table(tableName)}
3165
}
3266

3367
func listUsers(gid int64) (users, error) {

‎niu/main.go‎

Lines changed: 3 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ const (
2222
)
2323

2424
var (
25-
db *gorm.DB
25+
db *model
2626
globalLock sync.Mutex
2727

2828
errCancelFail = errors.New("遇到不可抗力因素，注销失败！")
@@ -84,27 +84,22 @@ func init() {
8484
panic(err)
8585
}
8686

87-
db = sdb.LogMode(false)
87+
db = &model{sdb.LogMode(false)}
8888

89+
registerTableHook(ensureUserInfo, ensureAuctionInfo)
8990
}
9091

9192
// DeleteWordNiuNiu ...
9293
func DeleteWordNiuNiu(gid, uid int64) error {
9394
globalLock.Lock()
9495
defer globalLock.Unlock()
95-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
96-
return err
97-
}
9896
return deleteUserByID(gid, uid)
9997
}
10098

10199
// SetWordNiuNiu length > 0 就增加 , length < 0 就减小
102100
func SetWordNiuNiu(gid, uid int64, length float64) error {
103101
globalLock.Lock()
104102
defer globalLock.Unlock()
105-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
106-
return err
107-
}
108103
m := map[string]interface{}{
109104
"length": length,
110105
}
@@ -116,10 +111,6 @@ func GetWordNiuNiu(gid, uid int64) (float64, error) {
116111
globalLock.Lock()
117112
defer globalLock.Unlock()
118113

119-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
120-
return 0, err
121-
}
122-
123114
info, err := getUserByID(gid, uid)
124115
if errors.Is(err, gorm.ErrRecordNotFound) {
125116
return 0, ErrNoNiuNiu
@@ -138,10 +129,6 @@ func GetRankingInfo(gid int64, t bool) (BaseInfos, error) {
138129
list users
139130
)
140131

141-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
142-
return nil, err
143-
}
144-
145132
us, err := listUsers(gid)
146133
if err != nil {
147134
return nil, err
@@ -170,9 +157,6 @@ func GetGroupUserRank(gid, uid int64) (int, error) {
170157
globalLock.Lock()
171158
defer globalLock.Unlock()
172159

173-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
174-
return 0, err
175-
}
176160
niu, err := getUserByID(gid, uid)
177161
if err != nil {
178162
return -1, err
@@ -188,9 +172,6 @@ func GetGroupUserRank(gid, uid int64) (int, error) {
188172

189173
// View 查看牛牛
190174
func View(gid, uid int64, name string) (string, error) {
191-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
192-
return "", err
193-
}
194175
i, err := getUserByID(gid, uid)
195176
if err != nil {
196177
return "", ErrNoNiuNiu
@@ -215,9 +196,6 @@ func View(gid, uid int64, name string) (string, error) {
215196

216197
// HitGlue 打胶
217198
func HitGlue(gid, uid int64, prop string) (string, error) {
218-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
219-
return "", err
220-
}
221199
niuniu, err := getUserByID(gid, uid)
222200
if err != nil {
223201
return "", ErrNoNiuNiuTwo
@@ -240,10 +218,6 @@ func Register(gid, uid int64) (string, error) {
240218
globalLock.Lock()
241219
defer globalLock.Unlock()
242220

243-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
244-
return "", err
245-
}
246-
247221
if _, err := getUserByID(gid, uid); err == nil {
248222
return "", ErrAlreadyRegistered
249223
}
@@ -273,10 +247,6 @@ func JJ(gid, uid, adduser int64, prop string) (message string, adduserLength flo
273247
globalLock.Lock()
274248
defer globalLock.Unlock()
275249

276-
if err = ensureUserInfoTable[userInfo](gid, ur); err != nil {
277-
return "", 0, ErrNoNiuNiu
278-
}
279-
280250
myniuniu, err := getUserByID(gid, uid)
281251
if err != nil {
282252
return "", 0, ErrNoNiuNiu
@@ -313,9 +283,6 @@ func JJ(gid, uid, adduser int64, prop string) (message string, adduserLength flo
313283
func Cancel(gid, uid int64) (string, error) {
314284
globalLock.Lock()
315285
defer globalLock.Unlock()
316-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
317-
return "", err
318-
}
319286
_, err := getUserByID(gid, uid)
320287
if err != nil {
321288
return "", ErrNoNiuNiuTwo
@@ -332,9 +299,6 @@ func Cancel(gid, uid int64) (string, error) {
332299
func Redeem(gid, uid int64, lastLength float64) error {
333300
globalLock.Lock()
334301
defer globalLock.Unlock()
335-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
336-
return err
337-
}
338302

339303
_, err := getUserByID(gid, uid)
340304
if err != nil {
@@ -369,9 +333,6 @@ func Redeem(gid, uid int64, lastLength float64) error {
369333
func Store(gid, uid int64, productID int, quantity int) error {
370334
globalLock.Lock()
371335
defer globalLock.Unlock()
372-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
373-
return err
374-
}
375336

376337
info, err := getUserByID(gid, uid)
377338
if err != nil {
@@ -398,12 +359,6 @@ func Store(gid, uid int64, productID int, quantity int) error {
398359
func Sell(gid, uid int64) (string, error) {
399360
globalLock.Lock()
400361
defer globalLock.Unlock()
401-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
402-
return "", err
403-
}
404-
if err := ensureUserInfoTable[AuctionInfo](gid, ac); err != nil {
405-
return "", err
406-
}
407362
niu, err := getUserByID(gid, uid)
408363
if err != nil {
409364
return "", ErrNoNiuNiu
@@ -444,19 +399,13 @@ func Sell(gid, uid int64) (string, error) {
444399
func ShowAuction(gid int64) ([]AuctionInfo, error) {
445400
globalLock.Lock()
446401
defer globalLock.Unlock()
447-
if err := ensureUserInfoTable[AuctionInfo](gid, ac); err != nil {
448-
return nil, err
449-
}
450402
return listAuction(gid)
451403
}
452404

453405
// Auction 购买牛牛
454406
func Auction(gid, uid int64, index int) (string, error) {
455407
globalLock.Lock()
456408
defer globalLock.Unlock()
457-
if err := ensureUserInfoTable[AuctionInfo](gid, ac); err != nil {
458-
return "", err
459-
}
460409

461410
infos, err := listAuction(gid)
462411
if len(infos) == 0 || err != nil {
@@ -511,9 +460,6 @@ func Auction(gid, uid int64, index int) (string, error) {
511460
func Bag(gid, uid int64) (string, error) {
512461
globalLock.Lock()
513462
defer globalLock.Unlock()
514-
if err := ensureUserInfoTable[userInfo](gid, ur); err != nil {
515-
return "", err
516-
}
517463
niu, err := getUserByID(gid, uid)
518464
if err != nil {
519465
return "", ErrNoNiuNiu

0 commit comments

Comments
 (0)