/
db.go
89 lines (73 loc) · 1.68 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
package middleware
import (
"context"
"github.com/kayx-org/freja/healthcheck"
"time"
)
//go:generate moq -out db_mock_test.go . db
type db interface {
Close() error
PingContext(context.Context) error
}
type OptionDbMiddleware func(*dbMiddleware)
type dbMiddleware struct {
db db
name string
checkWindow time.Duration
status healthcheck.ServiceStatus
}
// NewDB returns a new DB middleware which also implements the HealthCheck interface and can be configured accordinginly
func NewDB(db db, options ...OptionDbMiddleware) *dbMiddleware {
midDb := &dbMiddleware{
db: db,
name: "db",
checkWindow: time.Second,
status: healthcheck.UP,
}
for _, op := range options {
op(midDb)
}
return midDb
}
func OptionWindowCheck(t time.Duration) OptionDbMiddleware {
return func(m *dbMiddleware) {
m.checkWindow = t
}
}
func OptionHealthCheckName(t string) OptionDbMiddleware {
return func(m *dbMiddleware) {
m.name = t
}
}
func (m *dbMiddleware) Init() error {
return nil
}
func (m *dbMiddleware) Run(ctx context.Context) error {
ticker := time.NewTicker(m.checkWindow)
for {
select {
case <-ctx.Done():
return nil
case <-ticker.C:
m.runStatusCheck(ctx)
}
}
}
func (m *dbMiddleware) Stop(context.Context) error {
return m.db.Close()
}
func (m *dbMiddleware) Name() string {
return m.name
}
func (m *dbMiddleware) runStatusCheck(ctx context.Context) {
ctx, cancel := context.WithTimeout(ctx, m.checkWindow)
defer cancel()
if err := m.db.PingContext(ctx); err != nil {
m.status = healthcheck.DOWN
} else {
m.status = healthcheck.UP
}
}
func (m *dbMiddleware) Status() healthcheck.ServiceStatus {
return m.status
}