Skip to content

Commit 0794b54

Browse files
kvchandrew-farries
kvch
and
andrew-farries
authored
Add simple coordinator for executing DBActions (xataio#940)
This PR adds a coordinator to migrations package. This component is responsible for executing all registered actions only once in the correct order. I integrated it already into `Roll` for the complete phase of the migrations. --------- Co-authored-by: Andrew Farries <andyrb@gmail.com>
1 parent fb00928 commit 0794b54

15 files changed

Lines changed: 395 additions & 39 deletions

‎pkg/migrations/coordinator.go‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
3+
package migrations
4+
5+
import (
6+
"context"
7+
"fmt"
8+
"slices"
9+
)
10+
11+
// Coordinator is responsible for executing a series of database actions in a specific orderedActions.
12+
// It ensures that each action is executed only once, even if it is added multiple times.
13+
type Coordinator struct {
14+
actions map[string]DBAction
15+
orderedActions []string
16+
}
17+
18+
func NewCoordinator(actions []DBAction) *Coordinator {
19+
actionsMap := make(map[string]DBAction, len(actions))
20+
orderedActions := make([]string, 0)
21+
for _, action := range actions {
22+
if act, exists := actionsMap[action.ID()]; exists {
23+
idx := slices.Index(orderedActions, act.ID())
24+
orderedActions = slices.Delete(orderedActions, idx, idx+1)
25+
} else {
26+
actionsMap[action.ID()] = action
27+
}
28+
orderedActions = append(orderedActions, action.ID())
29+
}
30+
return &Coordinator{
31+
actions: actionsMap,
32+
orderedActions: orderedActions,
33+
}
34+
}
35+
36+
// Execute runs all actions in the order they were added to the coordinator.
37+
func (c *Coordinator) Execute(ctx context.Context) error {
38+
for _, id := range c.orderedActions {
39+
action, exists := c.actions[id]
40+
if !exists {
41+
return fmt.Errorf("action %s not found", id)
42+
}
43+
if err := action.Execute(ctx); err != nil {
44+
return fmt.Errorf("failed to execute action %s: %w", id, err)
45+
}
46+
}
47+
return nil
48+
}

‎pkg/migrations/coordinator_test.go‎

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
3+
package migrations
4+
5+
import (
6+
"testing"
7+
8+
"github.com/stretchr/testify/require"
9+
"github.com/xataio/pgroll/pkg/backfill"
10+
"github.com/xataio/pgroll/pkg/schema"
11+
)
12+
13+
func TestCoordinator(t *testing.T) {
14+
type testCase map[string]struct {
15+
actions []DBAction
16+
expectedOrder []string
17+
}
18+
19+
testCases := testCase{
20+
"empty": {
21+
actions: []DBAction{},
22+
expectedOrder: []string{},
23+
},
24+
"single action": {
25+
actions: []DBAction{
26+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"),
27+
},
28+
expectedOrder: []string{"rename_duplicated_t1_column1"},
29+
},
30+
"multiple actions with duplicates": {
31+
actions: []DBAction{
32+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"),
33+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column2"),
34+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"), // Duplicate
35+
},
36+
expectedOrder: []string{"rename_duplicated_t1_column2", "rename_duplicated_t1_column1"},
37+
},
38+
"multiple actions with mutiple duplicated for renaming": {
39+
actions: []DBAction{
40+
NewDropColumnAction(nil, "t1", "column1"),
41+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"),
42+
NewDropColumnAction(nil, "t1", "column2"),
43+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column2"),
44+
NewDropColumnAction(nil, "t1", "column3"),
45+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column3"),
46+
NewDropColumnAction(nil, "t1", "column1"),
47+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"),
48+
NewDropColumnAction(nil, "t1", "column2"),
49+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column2"),
50+
},
51+
expectedOrder: []string{
52+
"drop_column_t1_column3",
53+
"rename_duplicated_t1_column3",
54+
"drop_column_t1_column1",
55+
"rename_duplicated_t1_column1",
56+
"drop_column_t1_column2",
57+
"rename_duplicated_t1_column2",
58+
},
59+
},
60+
"add same column multiple times to same table": {
61+
actions: []DBAction{
62+
NewCreateTableAction(nil, "test_table", "", ""),
63+
NewAddColumnAction(nil, "t1", Column{Name: "column1"}, false),
64+
NewAddColumnAction(nil, "t1", Column{Name: "column1"}, false), // Duplicate
65+
},
66+
expectedOrder: []string{"create_table_test_table", "add_column_t1_column1"},
67+
},
68+
"create table multiple time with different statement with the same name": {
69+
actions: []DBAction{
70+
NewCreateTableAction(nil, "test_table", "", ""),
71+
NewCreateTableAction(nil, "test_table", "id INT", ""),
72+
NewCommentTableAction(nil, "test_table", ptr("This is a test table")),
73+
NewCommentColumnAction(nil, "test_table", "id", ptr("This is a test column")),
74+
},
75+
expectedOrder: []string{"create_table_test_table", "comment_table_test_table", "comment_column_test_table_id"},
76+
},
77+
"drop default value on column": {
78+
actions: []DBAction{
79+
NewCreateTableAction(nil, "test_table", "", ""),
80+
NewAddColumnAction(nil, "test_table1", Column{Name: "column1", Default: ptr("default_value")}, false),
81+
NewDropDefaultValueAction(nil, "test_table", "column1"),
82+
},
83+
expectedOrder: []string{"create_table_test_table", "add_column_test_table1_column1", "drop_default_test_table_column1"},
84+
},
85+
"alter column multiple times rollback phase": {
86+
actions: []DBAction{
87+
NewDropColumnAction(nil, "t1", "column1"),
88+
NewDropFunctionAction(nil,
89+
backfill.TriggerFunctionName("t1", "column1"),
90+
backfill.TriggerFunctionName("t1", "__pgroll_new_column1"),
91+
),
92+
NewDropColumnAction(nil, "t1", backfill.CNeedsBackfillColumn),
93+
NewDropColumnAction(nil, "t1", "column1"),
94+
NewDropFunctionAction(nil,
95+
backfill.TriggerFunctionName("t1", "column1"),
96+
backfill.TriggerFunctionName("t1", "__pgroll_new_column1"),
97+
),
98+
NewDropColumnAction(nil, "t1", backfill.CNeedsBackfillColumn),
99+
},
100+
expectedOrder: []string{
101+
"drop_column_t1_column1",
102+
"drop_function__pgroll_trigger_t1_column1__pgroll_trigger_t1___pgroll_new_column1",
103+
"drop_column_t1_" + backfill.CNeedsBackfillColumn,
104+
},
105+
},
106+
"alter column multiple times complete phase": {
107+
actions: []DBAction{
108+
NewAlterSequenceOwnerAction(nil, "t1", "column1", TemporaryName("column1")),
109+
NewDropColumnAction(nil, "t1", "column1"),
110+
NewDropFunctionAction(nil,
111+
backfill.TriggerFunctionName("t1", "column1"),
112+
backfill.TriggerFunctionName("t1", TemporaryName("column1")),
113+
),
114+
NewDropColumnAction(nil, "t1", backfill.CNeedsBackfillColumn),
115+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"),
116+
NewAlterSequenceOwnerAction(nil, "t1", "column1", TemporaryName("column1")),
117+
NewDropColumnAction(nil, "t1", "column1"),
118+
NewDropFunctionAction(nil,
119+
backfill.TriggerFunctionName("t1", "column1"),
120+
backfill.TriggerFunctionName("t1", TemporaryName("column1")),
121+
),
122+
NewDropColumnAction(nil, "t1", backfill.CNeedsBackfillColumn),
123+
NewRenameDuplicatedColumnAction(nil, &schema.Table{Name: "t1"}, "column1"),
124+
},
125+
expectedOrder: []string{
126+
"alter_sequence_owner_t1_column1_to__pgroll_new_column1",
127+
"drop_column_t1_column1",
128+
"drop_function__pgroll_trigger_t1_column1__pgroll_trigger_t1__pgroll_new_column1",
129+
"drop_column_t1_" + backfill.CNeedsBackfillColumn,
130+
"rename_duplicated_t1_column1",
131+
},
132+
},
133+
"create table and add multiple constraints": {
134+
actions: []DBAction{
135+
NewCreateTableAction(nil, "test_table", "", ""),
136+
NewAddColumnAction(nil, "column1", Column{}, false),
137+
NewAddColumnAction(nil, "column2", Column{}, false),
138+
NewCreateUniqueIndexConcurrentlyAction(nil, "public", "test_table", "my_idx", "column1"),
139+
NewCreateFKConstraintAction(nil, "test_table", "column2", []string{"other_column"}, nil, false, false, false),
140+
NewCreateCheckConstraintAction(nil, "test_table", "my_check", "column1 > 0", []string{"column1"}, false, false),
141+
},
142+
expectedOrder: []string{
143+
"create_table_test_table",
144+
"add_column_column1_",
145+
"add_column_column2_",
146+
"create_unique_index_concurrently_test_table_my_idx",
147+
"create_fk_constraint_test_table_column2",
148+
"create_check_constraint_test_table_my_check",
149+
},
150+
},
151+
}
152+
153+
for name, tc := range testCases {
154+
t.Run(name, func(t *testing.T) {
155+
coordinator := NewCoordinator(tc.actions)
156+
if len(coordinator.orderedActions) != len(tc.expectedOrder) {
157+
t.Fatalf("expected order length %d, got %d", len(tc.expectedOrder), len(coordinator.orderedActions))
158+
}
159+
160+
require.Equal(t, tc.expectedOrder, coordinator.orderedActions, "order of actions does not match expected")
161+
})
162+
}
163+
}

0 commit comments

Comments
 (0)