package/aqi
1
0
Fork 0
mirror of https://github.com/wonli/aqi.git synced 2025-10-06 17:16:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

add callback

Browse Source
This commit is contained in:
ideaa
2024-07-08 08:58:39 +08:00
parent 52d80f632f
commit 2c31cabc5b
4 changed files with 40 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
store/use.go
View File

@@ -1,5 +1,9 @@
package store package store
import "gorm.io/gorm"
type callback func(db *gorm.DB)
func DB(configKey string) *MySQLStore { func DB(configKey string) *MySQLStore {
return &MySQLStore{configKey: configKey} return &MySQLStore{configKey: configKey}
} }

12
store/use_mysql.go
View File

@@ -14,6 +14,9 @@ import (
type MySQLStore struct { type MySQLStore struct {
configKey string configKey string
callback callback
hasCallback bool
} }
func (m *MySQLStore) Config() *config.MySQL { func (m *MySQLStore) Config() *config.MySQL {
@@ -26,6 +29,11 @@ func (m *MySQLStore) Config() *config.MySQL {
return r return r
} }
func (m *MySQLStore) Callback(fn callback) {
m.callback = fn
m.hasCallback = true
}
func (m *MySQLStore) Use() *gorm.DB { func (m *MySQLStore) Use() *gorm.DB {
r := m.Config() r := m.Config()
if r == nil { if r == nil {
@@ -49,6 +57,10 @@ func (m *MySQLStore) Use() *gorm.DB {
return nil return nil
} }
if m.hasCallback {
m.callback(db)
}
sqlDB, err := db.DB() sqlDB, err := db.DB()
if err != nil { if err != nil {
logger.SugarLog.Error("Error pinging database", zap.String("error", err.Error())) logger.SugarLog.Error("Error pinging database", zap.String("error", err.Error()))

12
store/use_sqlite.go
View File

@@ -14,6 +14,9 @@ import (
type SQLiteStore struct { type SQLiteStore struct {
configKey string configKey string
callback callback
hasCallback bool
} }
func (m *SQLiteStore) Config() *config.Sqlite { func (m *SQLiteStore) Config() *config.Sqlite {
@@ -26,6 +29,11 @@ func (m *SQLiteStore) Config() *config.Sqlite {
return r return r
} }
func (m *SQLiteStore) Callback(fn callback) {
m.callback = fn
m.hasCallback = true
}
func (m *SQLiteStore) Use() *gorm.DB { func (m *SQLiteStore) Use() *gorm.DB {
r := m.Config() r := m.Config()
if r == nil { if r == nil {
@@ -45,6 +53,10 @@ func (m *SQLiteStore) Use() *gorm.DB {
return nil return nil
} }
if m.hasCallback {
m.callback(db)
}
sqlDB, err := db.DB() sqlDB, err := db.DB()
if err != nil { if err != nil {
logger.SugarLog.Error("Ping SQLite error", logger.SugarLog.Error("Ping SQLite error",

12
store/use_sqlserver.go
View File

@@ -13,6 +13,9 @@ import (
type SqlServerStore struct { type SqlServerStore struct {
configKey string configKey string
callback callback
hasCallback bool
} }
func (m *SqlServerStore) Config() *config.SqlServer { func (m *SqlServerStore) Config() *config.SqlServer {
@@ -25,6 +28,11 @@ func (m *SqlServerStore) Config() *config.SqlServer {
return r return r
} }
func (m *SqlServerStore) Callback(fn callback) {
m.callback = fn
m.hasCallback = true
}
func (m *SqlServerStore) Use() *gorm.DB { func (m *SqlServerStore) Use() *gorm.DB {
r := m.Config() r := m.Config()
if r == nil { if r == nil {
@@ -44,6 +52,10 @@ func (m *SqlServerStore) Use() *gorm.DB {
return nil return nil
} }
if m.hasCallback {
m.callback(db)
}
sqlDB, err := db.DB() sqlDB, err := db.DB()
if err != nil { if err != nil {
logger.SugarLog.Errorf("%s (ping)", err.Error()) logger.SugarLog.Errorf("%s (ping)", err.Error())