-
Notifications
You must be signed in to change notification settings - Fork 0
/
db.go
131 lines (107 loc) · 2.35 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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package store
import (
"fmt"
"path"
"strings"
"github.com/jmoiron/sqlx"
_ "github.com/mattn/go-sqlite3"
)
var (
db *DB
dbPath string
)
func init() {
dbPath = "."
var err error
if db = NewDB(); db == nil {
panic(fmt.Errorf("loading db: %s\n", err))
}
}
const KvSchema = `
CREATE TABLE IF NOT EXISTS %s (
k varchar(255) NOT NULL,
v BLOB,
PRIMARY KEY(k)
);
`
type DB struct{ *sqlx.DB }
func NewDB() (db *DB) {
if db == nil {
temp, err := sqlx.Open("sqlite3", path.Join(dbPath, "noye.db"))
if err != nil {
return nil
}
db = &DB{temp}
}
return
}
func fixScriptName(script string) string {
script = script + "_script"
script = strings.Replace(script, ".", "_", -1)
return script
}
func Get(table, key string) (string, error) {
table = fixScriptName(table)
return db.Get(table, key)
}
func Set(table, key, data string) (err error) {
table = fixScriptName(table)
return db.Set(table, key, data)
}
func (d *DB) Close() {
d.DB.DB.Close()
}
func (d *DB) Set(table, key, data string) (err error) {
if err := d.CheckTable(table, KvSchema); err != nil {
return err
}
tx, err := d.Beginx()
if err != nil {
return err
}
input := []byte(data)
// try update, this is awful
res, err := tx.Execv(fmt.Sprintf("UPDATE %s SET v = ? WHERE k = ?", table), input, key)
if err != nil {
return err
}
// it touched a row
if n, _ := res.RowsAffected(); n > 0 {
tx.Commit()
return nil
}
// try insert
_, err = tx.Execv(fmt.Sprintf("INSERT INTO %s (k, v) VALUES (?, ?);", table), key, input)
if err == nil {
tx.Commit()
}
return
}
func (d *DB) Get(table, key string) (string, error) {
err := d.CheckTable(table, KvSchema)
if err != nil {
return "", err
}
temp := map[string]interface{}{}
row := d.QueryRowx(fmt.Sprintf("SELECT v FROM %s WHERE k = ?", table), key)
if err := row.MapScan(temp); err != nil {
// log.Errorf("no rows for %s, %s: %s\n", table, key, err)
return "", err
}
if res, ok := temp["v"].(string); ok {
return res, nil
}
return "", fmt.Errorf("couldn't find '%s' on '%s'", key, table)
}
func (d *DB) CheckTable(table, schema string) error {
res, err := d.Execv("SELECT name FROM sqlite_master WHERE type = 'table' AND name = ?;", table)
if err != nil {
return err
}
if i, _ := res.RowsAffected(); i == 0 {
if _, err := d.Exec(fmt.Sprintf(schema, table)); err != nil {
return err
}
}
return nil
}