/
pre_db.go
80 lines (69 loc) · 1.67 KB
/
pre_db.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
package mongo
import (
"fmt"
"time"
"github.com/xiaoenai/tp-micro/model/redis"
"gopkg.in/mgo.v2"
)
// PreDB preset *DB
type PreDB struct {
*DB
preFuncs map[string]func() error
inited bool
}
// NewPreDB creates a unconnected *DB
func NewPreDB() *PreDB {
return &PreDB{
DB: &DB{
cacheableDBs: make(map[string]*CacheableDB),
},
preFuncs: make(map[string]func() error),
}
}
// Init initialize *DB.
func (p *PreDB) Init(mgoConfig *Config, redisConfig *redis.Config) (err error) {
var cache *redis.Client
if !mgoConfig.NoCache && redisConfig != nil {
cache, err = redis.NewClient(redisConfig)
if err != nil {
return err
}
}
// connect to mongodb
db, err := mgo.DialWithInfo(mgoConfig.Source())
if err != nil {
return err
}
p.DB.Session = db
p.DB.mgoConfig = mgoConfig
p.DB.Cache = cache
p.DB.redisConfig = redisConfig
for _, preFunc := range p.preFuncs {
if err = preFunc(); err != nil {
return err
}
}
p.inited = true
return nil
}
// RegCacheableDB registers a cacheable table.
func (p *PreDB) RegCacheableDB(ormStructPtr Cacheable, cacheExpiration time.Duration) (*CacheableDB, error) {
if p.inited {
return p.DB.RegCacheableDB(ormStructPtr, cacheExpiration)
}
tableName := ormStructPtr.TableName()
if _, ok := p.preFuncs[ormStructPtr.TableName()]; ok {
return nil, fmt.Errorf("re-register cacheable table: %s", tableName)
}
var cacheableDB = new(CacheableDB)
var preFunc = func() error {
_cacheableDB, err := p.DB.RegCacheableDB(ormStructPtr, cacheExpiration)
if err == nil {
*cacheableDB = *_cacheableDB
p.DB.cacheableDBs[tableName] = cacheableDB
}
return err
}
p.preFuncs[tableName] = preFunc
return cacheableDB, nil
}