/
branch.go
201 lines (186 loc) · 5.71 KB
/
branch.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
package flow
import (
"context"
)
// BranchCheckFunc checks the target and returns true if the branch should be selected.
type BranchCheckFunc[T Steper] func(context.Context, T) (bool, error)
// If adds a conditional branch to the workflow.
//
// If(someStep, func(ctx context.Context, someStep *SomeStep) (bool, error) {
// // branch condition here, true -> Then, false -> Else.
// // if error is returned, then fail the selected branch step.
// }).
// Then(thenStep).
// Else(elseStep)
func If[T Steper](target T, check BranchCheckFunc[T]) *IfBranch[T] {
return &IfBranch[T]{Target: target, BranchCheck: BranchCheck[T]{Check: check}}
}
// IfBranch adds target step, then and else step to workflow,
// and check the target step and determine which branch to go.
type IfBranch[T Steper] struct {
Target T // the target to check
BranchCheck BranchCheck[T]
ThenStep []Steper
ElseStep []Steper
Cond Condition // Cond is the When condition for both ThenStep and ElseStep, not target Step!
}
// Then adds steps to the Then branch.
func (i *IfBranch[T]) Then(th ...Steper) *IfBranch[T] {
i.ThenStep = append(i.ThenStep, th...)
return i
}
// Else adds steps to the Else branch.
func (i *IfBranch[T]) Else(el ...Steper) *IfBranch[T] {
i.ElseStep = append(i.ElseStep, el...)
return i
}
// When adds a condition to both Then and Else steps, not the Target!
// Default to DefaultCondition.
func (i *IfBranch[T]) When(cond Condition) *IfBranch[T] {
i.Cond = cond
return i
}
func (i *IfBranch[T]) isThen(isThen bool) Condition {
return func(ctx context.Context, ups map[Steper]StepResult) StepStatus {
if status := ConditionOrDefault(i.Cond)(ctx, ups); status != Running {
return status
}
if i.BranchCheck.OK == isThen {
return Running
}
return Skipped
}
}
func (i *IfBranch[T]) AddToWorkflow() map[Steper]*StepConfig {
return Steps().Merge(
Steps(i.Target).AfterStep(func(ctx context.Context, s Steper, err error) error {
i.BranchCheck.Do(ctx, i.Target)
return err
}),
Steps(i.ThenStep...).When(i.isThen(true)),
Steps(i.ElseStep...).When(i.isThen(false)),
Steps(append(append([]Steper{},
i.ThenStep...), i.ElseStep...,
)...).
DependsOn(i.Target).
BeforeStep(func(ctx context.Context, s Steper) (context.Context, error) {
if i.BranchCheck.Error != nil {
return ctx, i.BranchCheck.Error
}
return ctx, nil
}),
).AddToWorkflow()
}
// Switch adds a switch branch to the workflow.
//
// Switch(someStep).
// Case(case1, func(ctx context.Context, someStep *SomeStep) (bool, error) {
// // branch condition here, true to select this branch
// // error will fail the case
// }).
// Default(defaultStep), // the step to run if all case checks return false
// )
func Switch[T Steper](target T) *SwitchBranch[T] {
return &SwitchBranch[T]{Target: target, CasesToCheck: make(map[Steper]*BranchCheck[T])}
}
// SwitchBranch adds target step, cases and default step to workflow,
// and check the target step and determine which branch to go.
type SwitchBranch[T Steper] struct {
Target T
CasesToCheck map[Steper]*BranchCheck[T]
DefaultStep []Steper
Cond Condition
}
// BranchCheck represents a branch to be checked.
type BranchCheck[T Steper] struct {
Check BranchCheckFunc[T]
OK bool
Error error
}
func (bc *BranchCheck[T]) Do(ctx context.Context, target T) {
bc.OK, bc.Error = bc.Check(ctx, target)
}
// Case adds a case to the switch branch.
func (s *SwitchBranch[T]) Case(step Steper, check BranchCheckFunc[T]) *SwitchBranch[T] {
return s.Cases([]Steper{step}, check)
}
// Cases adds multiple cases to the switch branch.
// The check function will be executed for each case step.
func (s *SwitchBranch[T]) Cases(steps []Steper, check BranchCheckFunc[T]) *SwitchBranch[T] {
for _, step := range steps {
s.CasesToCheck[step] = &BranchCheck[T]{Check: check}
}
return s
}
// Default adds default step(s) to the switch branch.
func (s *SwitchBranch[T]) Default(step ...Steper) *SwitchBranch[T] {
s.DefaultStep = append(s.DefaultStep, step...)
return s
}
// When adds a condition to all case steps and default, not the Target!
func (s *SwitchBranch[T]) When(cond Condition) *SwitchBranch[T] {
s.Cond = cond
return s
}
func (s *SwitchBranch[T]) isCase(c Steper) func(ctx context.Context, ups map[Steper]StepResult) StepStatus {
return func(ctx context.Context, ups map[Steper]StepResult) StepStatus {
if status := ConditionOrDefault(s.Cond)(ctx, ups); status != Running {
return status
}
if check, ok := s.CasesToCheck[c]; ok {
check.Do(ctx, s.Target)
if check.OK {
return Running
}
}
return Skipped
}
}
func (s *SwitchBranch[T]) isDefault(ctx context.Context, ups map[Steper]StepResult) StepStatus {
for _, check := range s.CasesToCheck {
if check.OK {
return Skipped
}
}
// default branch ignores the status from cases
up := make(map[Steper]StepResult)
for step, status := range ups {
if _, isCase := s.CasesToCheck[step]; !isCase {
up[step] = status
}
}
if status := ConditionOrDefault(s.Cond)(ctx, up); status != Running {
return status
}
return Running
}
func (s *SwitchBranch[T]) AddToWorkflow() map[Steper]*StepConfig {
steps := Steps()
cases := []Steper{}
for step := range s.CasesToCheck {
step := step
cases = append(cases, step)
steps.Merge(
Steps(step).
DependsOn(s.Target).
When(s.isCase(step)).
BeforeStep(func(ctx context.Context, step Steper) (context.Context, error) {
for c, check := range s.CasesToCheck {
if HasStep(step, c) && check.Error != nil {
return ctx, check.Error
}
}
return ctx, nil
}),
)
}
if s.DefaultStep != nil {
steps.Merge(
Steps(s.DefaultStep...).
DependsOn(s.Target).
DependsOn(cases...).
When(s.isDefault),
)
}
return steps.AddToWorkflow()
}