/
database.go
127 lines (113 loc) · 2.91 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
package main
import (
"os"
"strings"
"time"
"github.com/jinzhu/gorm"
_ "github.com/mattn/go-sqlite3"
)
const (
sqlSelectByFreq = "cmd, dir, '' as hostname, '' as shell_session_id, time, count(*) c" // TODO: do max(time) here.
sqlGroupByFreq = "cmd"
sqlSelectByDate = "cmd, dir, hostname, shell_session_id, time"
)
// setting holds the persisted settings.
type setting struct {
ID int
SortByFreq bool
OnlyMySession bool
OnlyMyCwd bool
}
// record holds the data recorded for a single shell command.
type record struct {
Cmd string `sql:"size:65535"`
Dir string `sql:"size:65535"`
Hostname string `sql:"size:65535"`
ShellSessionID string `sql:"size:65535"`
Time time.Time
}
// query holds the components of a database query.
type query struct {
Cmd *string
Dir *string
Hostname *string
ShellSessionID *string
SortByFreq bool
Limit int
Offset int
}
// database holds a database connection and provides insert and retrieval.
type database interface {
Add(*record) error
Close() error
Query(query) ([]record, error)
Setting() (setting, error)
WriteSetting(*setting) error
}
type sqlDatabase struct {
db *gorm.DB
database
}
func newDatabase(path string) (database, error) {
d := &sqlDatabase{}
// Check if the DB already exists, or if we must create the table.
_, err := os.Stat(path)
n := os.IsNotExist(err)
d.db, err = gorm.Open("sqlite3", path)
// If new, we must create the table.
if n {
if err := d.db.CreateTable(&setting{}).Error; err != nil {
return nil, err
}
if err := d.db.CreateTable(&record{}).Error; err != nil {
return nil, err
}
}
return d, err
}
func (d *sqlDatabase) Add(r *record) error {
return d.db.Create(r).Error
}
func (d *sqlDatabase) Close() error {
return d.db.Close()
}
func (d *sqlDatabase) Query(q query) ([]record, error) {
var rs []record
var ws []string
var ps []interface{}
if q.Cmd != nil {
ws = append(ws, "cmd LIKE ?")
ps = append(ps, "%"+*q.Cmd+"%")
}
if q.Dir != nil {
ws = append(ws, "dir = ?")
ps = append(ps, q.Dir)
}
if q.Hostname != nil {
ws = append(ws, "hostname = ?")
ps = append(ps, q.Hostname)
}
if q.ShellSessionID != nil {
ws = append(ws, "shell_session_id = ?")
ps = append(ps, q.ShellSessionID)
}
var db *gorm.DB
if q.SortByFreq {
db = d.db.Table("records").Select(sqlSelectByFreq).Where(strings.Join(ws, " and "), ps...).Group(sqlGroupByFreq).Order("c desc").Limit(q.Limit).Offset(q.Offset)
} else {
db = d.db.Table("records").Select(sqlSelectByDate).Where(strings.Join(ws, " and "), ps...).Order("time desc").Limit(q.Limit).Offset(q.Offset)
}
err := db.Scan(&rs).Error
return rs, err
}
func (d *sqlDatabase) Setting() (setting, error) {
var s setting
r := d.db.First(&s)
if r.RecordNotFound() {
return s, nil
}
return s, r.Error
}
func (d *sqlDatabase) WriteSetting(s *setting) error {
return d.db.Save(s).Error
}