forked from go-pkgz/auth
/
bolt.go
129 lines (113 loc) · 4.05 KB
/
bolt.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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
package avatar
import (
"bytes"
"fmt"
"io"
"log"
"github.com/pkg/errors"
bolt "go.etcd.io/bbolt"
)
// BoltDB implements avatar store with bolt
// using separate db (file) with "avatars" bucket to keep image bin and "metas" bucket
// to keep sha1 of picture. avatarID (base file name) used as a key for both.
type BoltDB struct {
fileName string // full path to boltdb
db *bolt.DB
}
const avatarsBktName = "avatars"
const metasBktName = "metas"
// NewBoltDB makes bolt avatar store
func NewBoltDB(fileName string, options bolt.Options) (*BoltDB, error) {
db, err := bolt.Open(fileName, 0600, &options) //nolint
if err != nil {
return nil, errors.Wrapf(err, "failed to make boltdb for %s", fileName)
}
err = db.Update(func(tx *bolt.Tx) error {
if _, e := tx.CreateBucketIfNotExists([]byte(avatarsBktName)); e != nil {
return errors.Wrapf(e, "failed to create top level bucket %s", avatarsBktName)
}
_, e := tx.CreateBucketIfNotExists([]byte(metasBktName))
return errors.Wrapf(e, "failed to create top metas bucket %s", metasBktName)
})
if err != nil {
return nil, errors.Wrapf(err, "failed to initialize boltdb db %q buckets", fileName)
}
return &BoltDB{db: db, fileName: fileName}, nil
}
// Put avatar to bolt, key by avatarID. Trying to resize image and lso calculates sha1 of the file for ID func
func (b *BoltDB) Put(userID string, reader io.Reader) (avatar string, err error) {
id := encodeID(userID)
avatarID := id + imgSfx
err = b.db.Update(func(tx *bolt.Tx) error {
buf := &bytes.Buffer{}
if _, err = io.Copy(buf, reader); err != nil {
return errors.Wrapf(err, "can't read avatar %s", avatarID)
}
if err = tx.Bucket([]byte(avatarsBktName)).Put([]byte(avatarID), buf.Bytes()); err != nil {
return errors.Wrapf(err, "can't put to bucket with %s", avatarID)
}
// store sha1 of the image
return tx.Bucket([]byte(metasBktName)).Put([]byte(avatarID), []byte(hash(buf.Bytes(), avatarID)))
})
return avatarID, err
}
// Get avatar reader for avatar id.image, avatarID used as the direct key
func (b *BoltDB) Get(avatarID string) (reader io.ReadCloser, size int, err error) {
buf := &bytes.Buffer{}
err = b.db.View(func(tx *bolt.Tx) error {
data := tx.Bucket([]byte(avatarsBktName)).Get([]byte(avatarID))
if data == nil {
return fmt.Errorf("can't load avatar %s", avatarID)
}
size, err = buf.Write(data)
return errors.Wrapf(err, "failed to write for %s", avatarID)
})
return io.NopCloser(buf), size, err
}
// ID returns a fingerprint of the avatar content.
func (b *BoltDB) ID(avatarID string) (id string) {
data := []byte{}
err := b.db.View(func(tx *bolt.Tx) error {
if data = tx.Bucket([]byte(metasBktName)).Get([]byte(avatarID)); data == nil {
return fmt.Errorf("can't load avatar's id for %s", avatarID)
}
return nil
})
if err != nil { // failed to get ID, use encoded avatarID
log.Printf("[DEBUG] can't get avatar info '%s', %s", avatarID, err)
return encodeID(avatarID)
}
return string(data)
}
// Remove avatar from bolt
func (b *BoltDB) Remove(avatarID string) (err error) {
return b.db.Update(func(tx *bolt.Tx) error {
bkt := tx.Bucket([]byte(avatarsBktName))
if bkt.Get([]byte(avatarID)) == nil {
return errors.Errorf("avatar key not found, %s", avatarID)
}
if err = tx.Bucket([]byte(avatarsBktName)).Delete([]byte(avatarID)); err != nil {
return errors.Wrapf(err, "can't delete avatar object %s", avatarID)
}
return errors.Wrapf(tx.Bucket([]byte(metasBktName)).Delete([]byte(avatarID)),
"can't delete meta object %s", avatarID)
})
}
// List all avatars (ids) from metas bucket
// note: id includes .image suffix
func (b *BoltDB) List() (ids []string, err error) {
err = b.db.View(func(tx *bolt.Tx) error {
return tx.Bucket([]byte(metasBktName)).ForEach(func(k, _ []byte) error {
ids = append(ids, string(k))
return nil
})
})
return ids, fmt.Errorf("failed to list: %w", err)
}
// Close bolt store
func (b *BoltDB) Close() error {
return errors.Wrapf(b.db.Close(), "failed to close %s", b.fileName)
}
func (b *BoltDB) String() string {
return fmt.Sprintf("boltdb, path=%s", b.fileName)
}