|
| 1 | +begin; |
| 2 | +select plan(6); |
| 3 | +select pgflow_tests.reset_db(); |
| 4 | + |
| 5 | +-- ========================================================================= |
| 6 | +-- Test: Conditional flow_input in start_tasks |
| 7 | +-- |
| 8 | +-- Only root non-map steps receive flow_input. |
| 9 | +-- All other step types (root map, dependent non-map, dependent map) get NULL. |
| 10 | +-- This optimization prevents data duplication for large array processing. |
| 11 | +-- ========================================================================= |
| 12 | + |
| 13 | +-- Test 1: Root non-map step receives flow_input |
| 14 | +select pgflow.create_flow('root_step_flow'); |
| 15 | +select pgflow.add_step('root_step_flow', 'root_step'); |
| 16 | +select pgflow.start_flow('root_step_flow', '{"user_id": "abc123"}'::jsonb); |
| 17 | + |
| 18 | +select pgflow_tests.ensure_worker('root_step_flow'); |
| 19 | + |
| 20 | +with msgs as ( |
| 21 | + select * from pgmq.read_with_poll('root_step_flow', 10, 5, 1, 50) limit 1 |
| 22 | +), |
| 23 | +msg_ids as ( |
| 24 | + select array_agg(msg_id) as ids from msgs |
| 25 | +), |
| 26 | +started_tasks as ( |
| 27 | + select * from pgflow.start_tasks( |
| 28 | + 'root_step_flow', |
| 29 | + (select ids from msg_ids), |
| 30 | + '11111111-1111-1111-1111-111111111111'::uuid |
| 31 | + ) |
| 32 | +) |
| 33 | +select is( |
| 34 | + (select flow_input from started_tasks), |
| 35 | + '{"user_id": "abc123"}'::jsonb, |
| 36 | + 'Root non-map step should receive flow_input' |
| 37 | +); |
| 38 | + |
| 39 | +-- Test 2: Dependent non-map step receives NULL flow_input |
| 40 | +select pgflow_tests.reset_db(); |
| 41 | +select pgflow.create_flow('dep_flow'); |
| 42 | +select pgflow.add_step('dep_flow', 'first'); |
| 43 | +select pgflow.add_step('dep_flow', 'second', ARRAY['first']); |
| 44 | +select pgflow.start_flow('dep_flow', '{"original": "input"}'::jsonb); |
| 45 | + |
| 46 | +select pgflow_tests.ensure_worker('dep_flow'); |
| 47 | + |
| 48 | +-- Complete first step |
| 49 | +with poll_result as ( |
| 50 | + select * from pgflow_tests.read_and_start('dep_flow', 1, 1) |
| 51 | +) |
| 52 | +select pgflow.complete_task( |
| 53 | + run_id, |
| 54 | + step_slug, |
| 55 | + task_index, |
| 56 | + '{"first_result": "done"}'::jsonb |
| 57 | +) from poll_result; |
| 58 | + |
| 59 | +-- Start second step and verify flow_input is NULL |
| 60 | +select pgflow_tests.ensure_worker('dep_flow', '22222222-2222-2222-2222-222222222222'::uuid); |
| 61 | +with msgs as ( |
| 62 | + select * from pgmq.read_with_poll('dep_flow', 10, 5, 1, 50) limit 1 |
| 63 | +), |
| 64 | +msg_ids as ( |
| 65 | + select array_agg(msg_id) as ids from msgs |
| 66 | +), |
| 67 | +started_tasks as ( |
| 68 | + select * from pgflow.start_tasks( |
| 69 | + 'dep_flow', |
| 70 | + (select ids from msg_ids), |
| 71 | + '22222222-2222-2222-2222-222222222222'::uuid |
| 72 | + ) |
| 73 | +) |
| 74 | +select is( |
| 75 | + (select flow_input from started_tasks), |
| 76 | + NULL::jsonb, |
| 77 | + 'Dependent non-map step should receive NULL flow_input (lazy loaded)' |
| 78 | +); |
| 79 | + |
| 80 | +-- Test 3: Root map step receives NULL flow_input |
| 81 | +-- (flowInput IS the array, useless to include - workers get element via task.input) |
| 82 | +select pgflow_tests.reset_db(); |
| 83 | +select pgflow.create_flow('root_map_flow'); |
| 84 | +select pgflow.add_step('root_map_flow', 'map_step', '{}', null, null, null, null, 'map'); |
| 85 | +select pgflow.start_flow('root_map_flow', '[1, 2, 3]'::jsonb); |
| 86 | + |
| 87 | +select pgflow_tests.ensure_worker('root_map_flow'); |
| 88 | + |
| 89 | +with msgs as ( |
| 90 | + select * from pgmq.read_with_poll('root_map_flow', 10, 5, 3, 50) limit 1 |
| 91 | +), |
| 92 | +msg_ids as ( |
| 93 | + select array_agg(msg_id) as ids from msgs |
| 94 | +), |
| 95 | +started_tasks as ( |
| 96 | + select * from pgflow.start_tasks( |
| 97 | + 'root_map_flow', |
| 98 | + (select ids from msg_ids), |
| 99 | + '11111111-1111-1111-1111-111111111111'::uuid |
| 100 | + ) |
| 101 | +) |
| 102 | +select is( |
| 103 | + (select flow_input from started_tasks limit 1), |
| 104 | + NULL::jsonb, |
| 105 | + 'Root map step should receive NULL flow_input (flowInput is the array itself)' |
| 106 | +); |
| 107 | + |
| 108 | +-- Test 4: Dependent map step receives NULL flow_input |
| 109 | +select pgflow_tests.reset_db(); |
| 110 | +select pgflow.create_flow('dep_map_flow'); |
| 111 | +select pgflow.add_step('dep_map_flow', 'fetch_items', '{}', null, null, null, null, 'single'); |
| 112 | +select pgflow.add_step('dep_map_flow', 'process_item', ARRAY['fetch_items'], null, null, null, null, 'map'); |
| 113 | +select pgflow.start_flow('dep_map_flow', '{"config": "value"}'::jsonb); |
| 114 | + |
| 115 | +select pgflow_tests.ensure_worker('dep_map_flow'); |
| 116 | + |
| 117 | +-- Complete first step with array |
| 118 | +with poll_result as ( |
| 119 | + select * from pgflow_tests.read_and_start('dep_map_flow', 1, 1) |
| 120 | +) |
| 121 | +select pgflow.complete_task( |
| 122 | + run_id, |
| 123 | + step_slug, |
| 124 | + task_index, |
| 125 | + '["a", "b", "c"]'::jsonb |
| 126 | +) from poll_result; |
| 127 | + |
| 128 | +-- Start map step and verify flow_input is NULL |
| 129 | +select pgflow_tests.ensure_worker('dep_map_flow', '33333333-3333-3333-3333-333333333333'::uuid); |
| 130 | +with msgs as ( |
| 131 | + select * from pgmq.read_with_poll('dep_map_flow', 10, 5, 3, 50) limit 1 |
| 132 | +), |
| 133 | +msg_ids as ( |
| 134 | + select array_agg(msg_id) as ids from msgs |
| 135 | +), |
| 136 | +started_tasks as ( |
| 137 | + select * from pgflow.start_tasks( |
| 138 | + 'dep_map_flow', |
| 139 | + (select ids from msg_ids), |
| 140 | + '33333333-3333-3333-3333-333333333333'::uuid |
| 141 | + ) |
| 142 | +) |
| 143 | +select is( |
| 144 | + (select flow_input from started_tasks limit 1), |
| 145 | + NULL::jsonb, |
| 146 | + 'Dependent map step should receive NULL flow_input (lazy loaded)' |
| 147 | +); |
| 148 | + |
| 149 | +-- Test 5: Multiple parallel root non-map steps all receive flow_input |
| 150 | +select pgflow_tests.reset_db(); |
| 151 | +select pgflow.create_flow('parallel_flow'); |
| 152 | +select pgflow.add_step('parallel_flow', 'step1'); |
| 153 | +select pgflow.add_step('parallel_flow', 'step2'); -- parallel step, no deps |
| 154 | +select pgflow.start_flow('parallel_flow', '{"batch": "test"}'::jsonb); |
| 155 | + |
| 156 | +select pgflow_tests.ensure_worker('parallel_flow'); |
| 157 | + |
| 158 | +with msgs as ( |
| 159 | + select * from pgmq.read_with_poll('parallel_flow', 10, 5, 2, 50) |
| 160 | +), |
| 161 | +msg_ids as ( |
| 162 | + select array_agg(msg_id) as ids from msgs |
| 163 | +), |
| 164 | +started_tasks as ( |
| 165 | + select * from pgflow.start_tasks( |
| 166 | + 'parallel_flow', |
| 167 | + (select ids from msg_ids), |
| 168 | + '11111111-1111-1111-1111-111111111111'::uuid |
| 169 | + ) |
| 170 | +) |
| 171 | +select ok( |
| 172 | + (select bool_and(flow_input = '{"batch": "test"}'::jsonb) from started_tasks), |
| 173 | + 'All parallel root non-map steps should receive flow_input' |
| 174 | +); |
| 175 | + |
| 176 | +-- Test 6: Mixed batch - only root non-map steps receive flow_input |
| 177 | +-- This tests a complex scenario with different step types in same batch |
| 178 | +select pgflow_tests.reset_db(); |
| 179 | +select pgflow.create_flow('mixed_flow'); |
| 180 | +select pgflow.add_step('mixed_flow', 'root_single'); -- root non-map |
| 181 | +select pgflow.add_step('mixed_flow', 'root_array'); -- root non-map (array is just single that returns array) |
| 182 | +select pgflow.start_flow('mixed_flow', '{"data": "value"}'::jsonb); |
| 183 | + |
| 184 | +select pgflow_tests.ensure_worker('mixed_flow'); |
| 185 | + |
| 186 | +-- Both are root non-map steps, both should get flow_input |
| 187 | +with msgs as ( |
| 188 | + select * from pgmq.read_with_poll('mixed_flow', 10, 5, 2, 50) |
| 189 | +), |
| 190 | +msg_ids as ( |
| 191 | + select array_agg(msg_id) as ids from msgs |
| 192 | +), |
| 193 | +started_tasks as ( |
| 194 | + select * from pgflow.start_tasks( |
| 195 | + 'mixed_flow', |
| 196 | + (select ids from msg_ids), |
| 197 | + '11111111-1111-1111-1111-111111111111'::uuid |
| 198 | + ) |
| 199 | +) |
| 200 | +select is( |
| 201 | + (select count(*)::int from started_tasks where flow_input is not null), |
| 202 | + 2, |
| 203 | + 'Both root non-map steps should receive non-null flow_input' |
| 204 | +); |
| 205 | + |
| 206 | +select finish(); |
| 207 | +rollback; |
0 commit comments