/
ql.go
178 lines (166 loc) · 4.8 KB
/
ql.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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
// Package qlstore is a ql-based storage engine for the SCS session package.
//
// A working ql database is required, containing a sessions table with
// the definition:
//
// CREATE TABLE sessions (
// token string,
// data blob,
// expiry time
// )
// CREATE INDEX sessions_expiry_idx ON sessions (expiry);
//
// The qlstore package provides a background 'cleanup' goroutine to delete expired
// session data. This stops the database table from holding on to invalid sessions
// indefinitely and growing unnecessarily large.
//
// Usage:
//
// func main() {
// // Establish a database/sql pool
// db, err := sql.Open("ql", "example.db")
// if err != nil {
// log.Fatal(err)
// }
// defer db.Close()
//
// // Create a new QLStore instance using the existing database/sql pool,
// // with a cleanup interval of 5 minutes.
// engine := qlstore.New(db, 5*time.Minute)
//
// sessionManager := session.Manage(engine)
// http.ListenAndServe(":4000", sessionManager(http.DefaultServeMux))
// }
//
// The qlstore package is underpinned by the pq driver (github.com/cznic/ql/driver).
package qlstore
import (
"database/sql"
"log"
"time"
// Register ql driver with database/sql
_ "github.com/cznic/ql/driver"
)
// QLStore represents the currently configured session storage engine.
type QLStore struct {
*sql.DB
stopCleanup chan bool
}
// New returns a new QLStore instance.
//
// The cleanupInterval parameter controls how frequently expired session data
// is removed by the background cleanup goroutine. Setting it to 0 prevents
// the cleanup goroutine from running (i.e. expired sessions will not be removed).
func New(db *sql.DB, cleanupInterval time.Duration) *QLStore {
q := &QLStore{DB: db}
if cleanupInterval > 0 {
go q.startCleanup(cleanupInterval)
}
return q
}
func (q *QLStore) startCleanup(interval time.Duration) {
q.stopCleanup = make(chan bool)
ticker := time.NewTicker(interval)
for {
select {
case <-ticker.C:
err := q.deleteExpired()
if err != nil {
log.Println(err)
}
case <-q.stopCleanup:
ticker.Stop()
return
}
}
}
// Delete removes a session token and corresponding data from the QLStore instance.
func (q *QLStore) Delete(token string) error {
_, err := execTx(q.DB, "DELETE FROM sessions where token=$1", token)
return err
}
func (q *QLStore) deleteExpired() error {
_, err := execTx(q.DB, "DELETE FROM sessions WHERE expiry < now()")
return err
}
// Find returns the data for a given session token from the QLStore instance. If
// the session token is not found or is expired, the returned exists flag will
// be set to false.
func (q *QLStore) Find(token string) ([]byte, bool, error) {
var data []byte
query := "SELECT data FROM sessions WHERE token=$1 AND now()<expiry"
err := q.QueryRow(query, token).Scan(&data)
if err != nil {
if err == sql.ErrNoRows {
return nil, false, nil
}
return nil, false, err
}
return data, true, nil
}
// Save adds a session token and data to the QLStore instance with the given expiry time.
// If the session token already exists then the data and expiry time are updated.
func (q *QLStore) Save(token string, b []byte, expiry time.Time) error {
_, ok, _ := q.Find(token)
if ok {
_, err := execTx(q.DB, `
UPDATE sessions data=$2,expiry=$3 WHERE token=$1
`, token, b, expiry)
return err
}
_, err := execTx(q.DB, `
INSERT INTO sessions (token , data, expiry) VALUES ($1,$2,$3)
`, token, b, expiry)
return err
}
func execTx(db *sql.DB, query string, args ...interface{}) (sql.Result, error) {
tx, err := db.Begin()
if err != nil {
return nil, err
}
defer func() {
_ = tx.Commit()
}()
r, err := tx.Exec(query, args...)
return r, err
}
// StopCleanup terminates the background cleanup goroutine for the QLStore instance.
// It's rare to terminate this; generally QLStore instances and their cleanup
// goroutines are intended to be long-lived and run for the lifetime of your
// application.
//
// There may be occasions though when your use of the QLStore is transient. An
// example is creating a new QLStore instance in a test function. In this scenario,
// the cleanup goroutine (which will run forever) will prevent the QLStore object
// from being garbage collected even after the test function has finished. You
// can prevent this by manually calling StopCleanup.
//
// Example:
//
// func TestExample(t *testing.T) {
// db, err := sql.Open("ql", "test.db")
// if err != nil {
// t.Fatal(err)
// }
// defer db.Close()
//
// engine := pgstore.New(db, time.Second)
// defer engine.StopCleanup()
//
// // Run test...
// }
func (q *QLStore) StopCleanup() {
if q.stopCleanup != nil {
q.stopCleanup <- true
}
}
//Table provides SQL for creating a session table in ql database
func Table() string {
return `
CREATE TABLE sessions (
token string,
data blob,
expiry time
)
`
}