-
Notifications
You must be signed in to change notification settings - Fork 1.7k
/
orm.go
300 lines (260 loc) · 9.25 KB
/
orm.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
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
package pipeline
import (
"context"
"database/sql"
"fmt"
"strings"
"time"
"github.com/jmoiron/sqlx"
"github.com/pkg/errors"
uuid "github.com/satori/go.uuid"
"gorm.io/gorm"
"github.com/smartcontractkit/chainlink/core/services/postgres"
"github.com/smartcontractkit/chainlink/core/store/models"
)
var (
ErrNoSuchBridge = errors.New("no such bridge exists")
)
//go:generate mockery --name ORM --output ./mocks/ --case=underscore
type ORM interface {
CreateSpec(ctx context.Context, tx *gorm.DB, pipeline Pipeline, maxTaskTimeout models.Interval) (int32, error)
CreateRun(db *gorm.DB, run *Run) (err error)
StoreRun(db *sql.DB, run *Run) (restart bool, err error)
UpdateTaskRunResult(db *sql.DB, taskID uuid.UUID, result interface{}) (run Run, start bool, err error)
InsertFinishedRun(db *gorm.DB, run Run, trrs []TaskRunResult, saveSuccessfulTaskRuns bool) (runID int64, err error)
DeleteRunsOlderThan(threshold time.Duration) error
FindRun(id int64) (Run, error)
GetAllRuns() ([]Run, error)
GetUnfinishedRuns(now time.Time, fn func(run Run) error) error
DB() *gorm.DB
}
type orm struct {
db *gorm.DB
}
var _ ORM = (*orm)(nil)
func NewORM(db *gorm.DB) *orm {
return &orm{db}
}
// The tx argument must be an already started transaction.
func (o *orm) CreateSpec(ctx context.Context, tx *gorm.DB, pipeline Pipeline, maxTaskDuration models.Interval) (int32, error) {
spec := Spec{
DotDagSource: pipeline.Source,
MaxTaskDuration: maxTaskDuration,
}
err := tx.Create(&spec).Error
if err != nil {
return 0, err
}
return spec.ID, errors.WithStack(err)
}
func (o *orm) CreateRun(db *gorm.DB, run *Run) (err error) {
if run.CreatedAt.IsZero() {
return errors.New("run.CreatedAt must be set")
}
if err = db.Create(run).Error; err != nil {
return errors.Wrap(err, "error inserting pipeline_run")
}
return err
}
// StoreRun will persist a partially executed run before suspending, or finish a run.
// If `restart` is true, then new task run data is available and the run should be resumed immediately.
func (o *orm) StoreRun(db *sql.DB, run *Run) (restart bool, err error) {
finished := run.FinishedAt.Valid
err = postgres.SqlxTransaction(context.Background(), db, func(tx *sqlx.Tx) error {
if !finished {
// Lock the current run. This prevents races with /v2/resume
sql := `SELECT id FROM pipeline_runs WHERE id = $1 FOR UPDATE;`
if _, err = tx.Exec(sql, run.ID); err != nil {
return err
}
taskRuns := []TaskRun{}
// Reload task runs, we want to check for any changes while the run was ongoing
if err = tx.Select(&taskRuns, `SELECT * FROM pipeline_task_runs WHERE pipeline_run_id = $1`, run.ID); err != nil {
return err
}
// Construct a temporary run so we can use r.ByDotID
tempRun := Run{PipelineTaskRuns: taskRuns}
// Diff with current state, if updated, swap run.PipelineTaskRuns and early return with restart = true
for i, tr := range run.PipelineTaskRuns {
if !tr.IsPending() {
continue
}
// Look for new data
if taskRun := tempRun.ByDotID(tr.DotID); taskRun != nil {
// Swap in the latest state
run.PipelineTaskRuns[i] = *taskRun
restart = true
}
}
if restart {
return nil
}
// Suspend the run
run.State = RunStatusSuspended
if _, err = tx.NamedExec(`UPDATE pipeline_runs SET state = :state`, run); err != nil {
return err
}
} else {
// Simply finish the run, no need to do any sort of locking
if run.Outputs.Val == nil || len(run.Errors) == 0 {
return errors.Errorf("run must have both Outputs and Errors, got Outputs: %#v, Errors: %#v", run.Outputs.Val, run.Errors)
}
sql := `UPDATE pipeline_runs SET state = :state, finished_at = :finished_at, errors= :errors, outputs = :outputs WHERE id = :id`
if _, err = tx.NamedExec(sql, run); err != nil {
return err
}
}
sql := `
INSERT INTO pipeline_task_runs (pipeline_run_id, id, type, index, output, error, dot_id, created_at, finished_at)
VALUES (:pipeline_run_id, :id, :type, :index, :output, :error, :dot_id, :created_at, :finished_at)
ON CONFLICT (pipeline_run_id, dot_id) DO UPDATE SET
output = EXCLUDED.output, error = EXCLUDED.error, finished_at = EXCLUDED.finished_at
RETURNING *;
`
// NOTE: can't use Select() to auto scan because we're using NamedQuery,
// sqlx.Named + Select is possible but it's about the same amount of code
var rows *sqlx.Rows
rows, err = tx.NamedQuery(sql, run.PipelineTaskRuns)
if err != nil {
return err
}
taskRuns := []TaskRun{}
if err = sqlx.StructScan(rows, &taskRuns); err != nil {
return err
}
// replace with new task run data
run.PipelineTaskRuns = taskRuns
return nil
})
return restart, err
}
func (o *orm) UpdateTaskRunResult(db *sql.DB, taskID uuid.UUID, result interface{}) (run Run, start bool, err error) {
err = postgres.SqlxTransaction(context.Background(), db, func(tx *sqlx.Tx) error {
sql := `
SELECT pipeline_runs.*, pipeline_specs.dot_dag_source "pipeline_spec.dot_dag_source"
FROM pipeline_runs
JOIN pipeline_task_runs ON (pipeline_task_runs.pipeline_run_id = pipeline_runs.id)
JOIN pipeline_specs ON (pipeline_specs.id = pipeline_runs.pipeline_spec_id)
WHERE pipeline_task_runs.id = $1 AND pipeline_runs.state in ('running', 'suspended')
FOR UPDATE`
if err = tx.Get(&run, sql, taskID); err != nil {
return err
}
// Update the task with result
sql = `UPDATE pipeline_task_runs SET output = $2, finished_at = $3 WHERE id = $1`
if _, err = tx.Exec(sql, taskID, JSONSerializable{Val: result}, time.Now()); err != nil {
return err
}
if run.State == RunStatusSuspended {
start = true
run.State = RunStatusRunning
// We're going to restart the run, so set it back to "in progress"
sql = `UPDATE pipeline_runs SET state = $2 WHERE id = $1`
if _, err = tx.Exec(sql, run.ID, run.State); err != nil {
return err
}
// NOTE: can't join and preload in a single query unless explicitly listing all the struct fields...
// https://snippets.aktagon.com/snippets/757-how-to-join-two-tables-with-jmoiron-sqlx
sql = `SELECT * FROM pipeline_task_runs WHERE pipeline_run_id = $1`
return tx.Select(&run.PipelineTaskRuns, sql, run.ID)
}
return nil
})
return run, start, err
}
// If saveSuccessfulTaskRuns = false, we only save errored runs.
// That way if the job is run frequently (such as OCR) we avoid saving a large number of successful task runs
// which do not provide much value.
func (o *orm) InsertFinishedRun(db *gorm.DB, run Run, trrs []TaskRunResult, saveSuccessfulTaskRuns bool) (runID int64, err error) {
if run.CreatedAt.IsZero() {
return 0, errors.New("run.CreatedAt must be set")
}
if run.FinishedAt.IsZero() {
return 0, errors.New("run.FinishedAt must be set")
}
if run.Outputs.Val == nil || len(run.Errors) == 0 {
return 0, errors.Errorf("run must have both Outputs and Errors, got Outputs: %#v, Errors: %#v", run.Outputs.Val, run.Errors)
}
if len(trrs) == 0 && (saveSuccessfulTaskRuns || run.HasErrors()) {
return 0, errors.New("must provide task run results")
}
err = postgres.GormTransactionWithoutContext(db, func(tx *gorm.DB) error {
if err = tx.Create(&run).Error; err != nil {
return errors.Wrap(err, "error inserting finished pipeline_run")
}
if !saveSuccessfulTaskRuns && !run.HasErrors() {
return nil
}
sql := `
INSERT INTO pipeline_task_runs (pipeline_run_id, id, type, index, output, error, dot_id, created_at, finished_at)
VALUES %s
`
valueStrings := []string{}
valueArgs := []interface{}{}
for _, trr := range trrs {
valueStrings = append(valueStrings, "(?,?,?,?,?,?,?,?,?)")
valueArgs = append(valueArgs, run.ID, trr.ID, trr.Task.Type(), trr.Task.OutputIndex(), trr.Result.OutputDB(), trr.Result.ErrorDB(), trr.Task.DotID(), trr.CreatedAt, trr.FinishedAt)
}
/* #nosec G201 */
stmt := fmt.Sprintf(sql, strings.Join(valueStrings, ","))
return tx.Exec(stmt, valueArgs...).Error
})
return run.ID, err
}
func (o *orm) DeleteRunsOlderThan(threshold time.Duration) error {
err := o.db.Exec(
`DELETE FROM pipeline_runs WHERE finished_at < ?`, time.Now().Add(-threshold),
).Error
if err != nil {
return err
}
return nil
}
func (o *orm) FindRun(id int64) (Run, error) {
var run = Run{ID: id}
err := o.db.
Preload("PipelineSpec").
Preload("PipelineTaskRuns", func(db *gorm.DB) *gorm.DB {
return db.
Order("created_at ASC, id ASC")
}).First(&run).Error
return run, err
}
func (o *orm) GetAllRuns() ([]Run, error) {
var runs []Run
err := o.db.
Preload("PipelineSpec").
Preload("PipelineTaskRuns", func(db *gorm.DB) *gorm.DB {
return db.
Order("created_at ASC, id ASC")
}).Find(&runs).Error
return runs, err
}
func (o *orm) GetUnfinishedRuns(now time.Time, fn func(run Run) error) error {
return postgres.Batch(func(offset, limit uint) (count uint, err error) {
var runs []Run
err = o.db.
Preload("PipelineSpec").
Preload("PipelineTaskRuns", func(db *gorm.DB) *gorm.DB {
return db.
Order("created_at ASC, id ASC")
}).
Where(`state = ? AND created_at < ?`, RunStatusRunning, now).
Order("created_at ASC, id ASC").
Limit(int(limit)).
Offset(int(offset)).
Find(&runs).Error
if err != nil {
return 0, err
}
for _, run := range runs {
if err := fn(run); err != nil {
return 0, err
}
}
return uint(len(runs)), nil
})
}
func (o *orm) DB() *gorm.DB {
return o.db
}