/
db.go
137 lines (111 loc) · 3.58 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
132
133
134
135
136
137
package account
import (
"encoding/json"
"errors"
"fmt"
"berty.tech/core/pkg/tracing"
"github.com/jinzhu/gorm"
opentracing "github.com/opentracing/opentracing-go"
)
//
// state db
//
type StateDB struct {
Gorm *gorm.DB `gorm:"-"`
gorm.Model
StartCounter int
JSONNetConf string
BotMode bool
LocalGRPC bool
}
func OpenStateDB(path string, initialState StateDB) (*StateDB, error) {
// open db
db, err := gorm.Open("sqlite3", path)
if err != nil {
return nil, err
}
// create tables and columns
if err := db.AutoMigrate(&StateDB{}).Error; err != nil {
return nil, err
}
// preload last state
var state StateDB
if err := db.Last(&state).Error; err != nil && err != gorm.ErrRecordNotFound {
return nil, err
} else if err == gorm.ErrRecordNotFound {
state = initialState
if err := db.Save(&state).Error; err != nil {
return nil, err
}
}
state.Gorm = db
return &state, nil
}
func (state StateDB) String() string {
out, _ := json.Marshal(state)
return string(out)
}
func (state *StateDB) Save() error {
return state.Gorm.Save(state).Error
}
func (state *StateDB) Create() error {
return state.Gorm.Create(state).Error
}
func (state *StateDB) Close() {
state.Gorm.Close()
}
//
// app db
//
func gormCreateSubSpan(scope *gorm.Scope, operationName string) {
if span, ok := scope.Get("rootSpan"); ok {
rootSpan := span.(opentracing.Span)
operationName = fmt.Sprintf("gorm::%s", operationName)
subSpan := rootSpan.Tracer().StartSpan(
operationName,
opentracing.ChildOf(rootSpan.Context()),
)
subSpan.SetTag("component", "gorm")
scope.Set("subSpan", subSpan)
}
}
func gormFinishSubSpan(scope *gorm.Scope) {
if span, ok := scope.Get("subSpan"); ok {
subSpan := span.(opentracing.Span)
defer subSpan.Finish()
subSpan.LogKV("sql", scope.SQL)
}
}
func WithDatabase(opts *DatabaseOptions) NewOption {
return func(a *Account) error {
tracer := tracing.EnterFunc(a.rootContext, opts)
defer tracer.Finish()
ctx := tracer.Context()
if opts == nil {
opts = &DatabaseOptions{}
}
a.dbDir = opts.Path
if a.dbDir == "" {
return errors.New("cannot have empty database path")
}
a.dbDrop = opts.Drop
if err := a.openDatabase(ctx); err != nil {
return err
}
if a.tracer != nil {
// create
a.db.Callback().Create().Before("gorm:before_create").Register("jaeger:before_create", func(scope *gorm.Scope) { gormCreateSubSpan(scope, fmt.Sprintln("insert", scope.TableName())) })
a.db.Callback().Create().Before("gorm:after_create").Register("jaeger:after_create", func(scope *gorm.Scope) { gormFinishSubSpan(scope) })
// update
a.db.Callback().Update().Before("gorm:before_update").Register("jaeger:before_update", func(scope *gorm.Scope) { gormCreateSubSpan(scope, fmt.Sprintln("update", scope.TableName())) })
a.db.Callback().Update().Before("gorm:after_update").Register("jaeger:after_update", func(scope *gorm.Scope) { gormFinishSubSpan(scope) })
// delete
a.db.Callback().Delete().Before("gorm:before_delete").Register("jaeger:before_delete", func(scope *gorm.Scope) { gormCreateSubSpan(scope, fmt.Sprintln("delete", scope.TableName())) })
a.db.Callback().Delete().Before("gorm:after_delete").Register("jaeger:after_delete", func(scope *gorm.Scope) { gormFinishSubSpan(scope) })
// query
a.db.Callback().Query().Before("gorm:query").Register("jaeger:before_query", func(scope *gorm.Scope) { gormCreateSubSpan(scope, fmt.Sprintln("select", scope.TableName())) })
a.db.Callback().Query().Before("gorm:after_query").Register("jaeger:after_query", func(scope *gorm.Scope) { gormFinishSubSpan(scope) })
}
return nil
}
}