-
Notifications
You must be signed in to change notification settings - Fork 2
/
database.go
139 lines (128 loc) · 3.06 KB
/
database.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
130
131
132
133
134
135
136
137
138
139
package main
import (
"fmt"
"github.com/boltdb/bolt"
"log"
"math/rand"
)
const alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
var db *bolt.DB
//https://zupzup.org/boltdb-example/
// Loads the database; Creates one if does not exist
func LoadDB(dataBaseName string) error {
var err error
db, err = bolt.Open(dataBaseName, 0600, nil)
if err != nil {
return fmt.Errorf("could not open db, %v", err)
}
err = db.Update(func(tx *bolt.Tx) error {
_, err := tx.CreateBucketIfNotExists([]byte("DB"))
if err != nil {
return fmt.Errorf("could not create root bucket: %v", err)
}
return nil
})
if err != nil {
return fmt.Errorf("could not set up buckets, %v", err)
}
log.Println("DB Setup Done")
return nil
}
func CloseDB() {
_ = db.Close()
}
//Inserts a link into the database
//Returns the ID for the link to be shared later
func InsertValue(Value string) (string, error) {
for {
key := generateRandomStringAsByte()
hasValue := false
//Check if the random exists in database
err := db.View(func(tx *bolt.Tx) error {
check := tx.Bucket([]byte("DB")).Get(key)
hasValue = check != nil
return nil
})
if err != nil {
return "", err
}
if !hasValue { //In this case we save the string into database
err = db.Update(func(tx *bolt.Tx) error {
err = tx.Bucket([]byte("DB")).Put(key, []byte(Value))
if err != nil {
return fmt.Errorf("could not read db: %v", err)
}
return nil
})
if err != nil {
return "", err
}
return string(key), nil
}
}
}
//Check if a key exists; On errors return false as well
func HasKey(Key string) bool {
hasValue := false
//Check if the random exists in database
_ = db.View(func(tx *bolt.Tx) error {
check := tx.Bucket([]byte("DB")).Get([]byte(Key))
hasValue = check != nil
return nil
})
return hasValue
}
//Remove a key from the database
func RemoveKey(Key string) error {
if !HasKey(Key) {
return fmt.Errorf("this token does not exits")
}
err := db.Update(func(tx *bolt.Tx) error {
err := tx.Bucket([]byte("DB")).Delete([]byte(Key))
if err != nil {
return fmt.Errorf("could not delete key: %v", err)
}
return nil
})
return err
}
//List all of the values
func ListAllValues() (map[string]string, error) {
m := make(map[string]string)
err := db.View(func(tx *bolt.Tx) error {
err := tx.Bucket([]byte("DB")).ForEach(func(k, v []byte) error {
if len(v) > 100 {
m[string(k)] = escapeMarkdown(string(v[:100])) + " *...* "
} else {
m[string(k)] = escapeMarkdown(string(v))
}
return nil
})
return err
})
return m, err
}
//Read the value from database
func ReadValue(Key string) (string, error) {
var res string
err := db.View(func(tx *bolt.Tx) error {
check := tx.Bucket([]byte("DB")).Get([]byte(Key))
if check == nil {
return fmt.Errorf("Cannot find value for " + Key)
}
res = string(check)
return nil
})
if err != nil {
return "", err
}
return res, nil
}
//Keys are 8 letter long
func generateRandomStringAsByte() []byte {
s := ""
for i := 0; i < 8; i++ {
s += string(alphabet[rand.Int31n(52)])
}
return []byte(s)
}