-
Notifications
You must be signed in to change notification settings - Fork 0
/
db.go
113 lines (91 loc) · 2.71 KB
/
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
package main
import (
"bytes"
"time"
bolt "github.com/coreos/bbolt"
"github.com/pkg/errors"
)
// DB is a cache, wraps bolt.DB
type DB struct {
ManagerBucketName []byte
ChannelBucketName []byte
*bolt.DB
}
// DefaultDBTimeout for Bolt
const DefaultDBTimeout = 1 * time.Second
func openDB(path *string, mBucket, cBucket string, timeout time.Duration) (*DB, error) {
if timeout == 0 {
timeout = DefaultDBTimeout
}
// open DB
boltDB, err := bolt.Open(*path, 0600, &bolt.Options{Timeout: timeout})
if err != nil {
return nil, err
}
db := &DB{[]byte(mBucket), []byte(cBucket), boltDB}
// create buckets if needed
if err = db.newBucket(db.ManagerBucketName); err != nil {
return nil, err
}
if err = db.newBucket(db.ChannelBucketName); err != nil {
return nil, err
}
return db, nil
}
// SaveUserBDToCache saves the record about the user's birthday into the cache DB's bucket
func (db *DB) SaveUserBDToCache(bucketName []byte, id, bd string) error {
if !bytes.Equal(bucketName, db.ManagerBucketName) && !bytes.Equal(bucketName, db.ChannelBucketName) {
return errors.Errorf("bucket %q does not exist", bucketName)
}
if err := db.put(bucketName, []byte(id), []byte(bd)); err != nil {
return errors.Wrap(err, "unable to put value into DB")
}
return nil
}
// CheckUserBDInCache checks if the record about the user's birthday is present in the cache DB's bucket
func (db *DB) CheckUserBDInCache(bucketName []byte, id, bd string) (bool, error) {
if !bytes.Equal(bucketName, db.ManagerBucketName) && !bytes.Equal(bucketName, db.ChannelBucketName) {
return false, errors.Errorf("bucket %q does not exist", bucketName)
}
dbValue, err := db.get(bucketName, []byte(id))
if err != nil {
return false, errors.Wrap(err, "unable to get value from DB")
}
if dbValue != nil {
// check if stored BD record is the same
if bytes.Equal([]byte(bd), dbValue) {
return true, nil
}
}
return false, nil
}
func (db *DB) newBucket(bucketName []byte) error {
return db.Update(func(tx *bolt.Tx) error {
_, err := tx.CreateBucketIfNotExists(bucketName)
return err
})
}
func (db *DB) put(bucketName, key, value []byte) error {
return db.Update(func(tx *bolt.Tx) error {
bucket := tx.Bucket(bucketName)
if bucket == nil {
return errors.Errorf("bucket %q not found", bucketName)
}
return bucket.Put(key, value)
})
}
func (db *DB) get(bucketName, key []byte) (value []byte, err error) {
err = db.View(func(tx *bolt.Tx) error {
bucket := tx.Bucket(bucketName)
if bucket == nil {
return errors.Errorf("bucket %q not found", bucketName)
}
if v := bucket.Get(key); v != nil {
// copy v into value as v only lives till the end of tx
value = make([]byte, len(v))
copy(value, v)
}
return nil
})
return
}