pgevolve_core/plan/
grouping.rs

1//! [`group_steps`] — partition a flat `Vec<RawStep>` into `TransactionGroup`s.
2//!
3//! Adjacent steps with the same [`TransactionConstraint`] coalesce into one
4//! group. The resulting list is the unit the executor consumes:
5//!
6//! - Transactional groups run inside a single `BEGIN; ... COMMIT;`.
7//! - Non-transactional groups are conceptually a sequence of singletons; each
8//!   step runs autocommit. Coalescing them into a "group" is purely
9//!   organizational — it keeps the directive-comment structure clean (spec §7.1).
10
11use crate::plan::raw_step::{RawStep, TransactionConstraint};
12
13/// A run of consecutive [`RawStep`]s that share a transactional / non-transactional
14/// classification. Group ids are 1-indexed in emission order.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct TransactionGroup {
17    /// 1-indexed group id (assigned by [`group_steps`]).
18    pub id: u32,
19    /// True if every step in the group can run inside a single
20    /// `BEGIN; ... COMMIT;`. Non-transactional groups (`CONCURRENTLY`, etc.)
21    /// run as autocommit singletons even when the group has multiple steps.
22    pub transactional: bool,
23    /// The steps in the group, in emission order.
24    pub steps: Vec<RawStep>,
25}
26
27/// Partition `steps` into groups by transactional boundary.
28///
29/// Adjacent steps with the same [`TransactionConstraint`] go in one group;
30/// each transition between in-tx and out-of-tx starts a new group. Empty
31/// input produces an empty output. Group ids are 1-indexed.
32pub fn group_steps(steps: Vec<RawStep>) -> Vec<TransactionGroup> {
33    let mut groups: Vec<TransactionGroup> = Vec::new();
34    let mut current: Vec<RawStep> = Vec::new();
35    let mut current_kind: Option<TransactionConstraint> = None;
36
37    for step in steps {
38        match current_kind {
39            None => {
40                current_kind = Some(step.transactional);
41                current.push(step);
42            }
43            Some(prev) if prev == step.transactional => {
44                current.push(step);
45            }
46            Some(prev) => {
47                let id = u32::try_from(groups.len() + 1).unwrap_or(u32::MAX);
48                groups.push(TransactionGroup {
49                    id,
50                    transactional: matches!(prev, TransactionConstraint::InTransaction),
51                    steps: std::mem::take(&mut current),
52                });
53                current_kind = Some(step.transactional);
54                current.push(step);
55            }
56        }
57    }
58
59    if let Some(prev) = current_kind {
60        let id = u32::try_from(groups.len() + 1).unwrap_or(u32::MAX);
61        groups.push(TransactionGroup {
62            id,
63            transactional: matches!(prev, TransactionConstraint::InTransaction),
64            steps: current,
65        });
66    }
67
68    groups
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::plan::raw_step::StepKind;
75
76    fn step(kind: StepKind, c: TransactionConstraint) -> RawStep {
77        RawStep {
78            step_no: 0,
79            kind,
80            destructive: false,
81            destructive_reason: None,
82            intent_id: None,
83            targets: vec![],
84            sql: String::new(),
85            transactional: c,
86        }
87    }
88
89    #[test]
90    fn empty_input_yields_no_groups() {
91        let out = group_steps(vec![]);
92        assert!(out.is_empty());
93    }
94
95    #[test]
96    fn single_in_tx_step_yields_one_group() {
97        let out = group_steps(vec![step(
98            StepKind::CreateTable,
99            TransactionConstraint::InTransaction,
100        )]);
101        assert_eq!(out.len(), 1);
102        assert_eq!(out[0].id, 1);
103        assert!(out[0].transactional);
104        assert_eq!(out[0].steps.len(), 1);
105    }
106
107    #[test]
108    fn single_out_of_tx_step_yields_one_group() {
109        let out = group_steps(vec![step(
110            StepKind::CreateIndexConcurrent,
111            TransactionConstraint::OutsideTransaction,
112        )]);
113        assert_eq!(out.len(), 1);
114        assert_eq!(out[0].id, 1);
115        assert!(!out[0].transactional);
116    }
117
118    #[test]
119    fn all_in_tx_steps_coalesce_into_one_group() {
120        let out = group_steps(vec![
121            step(StepKind::CreateSchema, TransactionConstraint::InTransaction),
122            step(StepKind::CreateTable, TransactionConstraint::InTransaction),
123            step(StepKind::CreateIndex, TransactionConstraint::InTransaction),
124        ]);
125        assert_eq!(out.len(), 1);
126        assert_eq!(out[0].steps.len(), 3);
127        assert!(out[0].transactional);
128    }
129
130    #[test]
131    fn transition_creates_new_group() {
132        let out = group_steps(vec![
133            step(StepKind::CreateTable, TransactionConstraint::InTransaction),
134            step(
135                StepKind::CreateIndexConcurrent,
136                TransactionConstraint::OutsideTransaction,
137            ),
138            step(
139                StepKind::ValidateConstraint,
140                TransactionConstraint::InTransaction,
141            ),
142        ]);
143        assert_eq!(out.len(), 3);
144        assert_eq!(out[0].id, 1);
145        assert!(out[0].transactional);
146        assert_eq!(out[1].id, 2);
147        assert!(!out[1].transactional);
148        assert_eq!(out[2].id, 3);
149        assert!(out[2].transactional);
150    }
151
152    #[test]
153    fn consecutive_out_of_tx_steps_coalesce() {
154        let out = group_steps(vec![
155            step(
156                StepKind::CreateIndexConcurrent,
157                TransactionConstraint::OutsideTransaction,
158            ),
159            step(
160                StepKind::DropIndexConcurrent,
161                TransactionConstraint::OutsideTransaction,
162            ),
163        ]);
164        assert_eq!(out.len(), 1);
165        assert!(!out[0].transactional);
166        assert_eq!(out[0].steps.len(), 2);
167    }
168
169    #[test]
170    fn group_ids_are_one_indexed_and_sequential() {
171        let out = group_steps(vec![
172            step(StepKind::CreateTable, TransactionConstraint::InTransaction),
173            step(
174                StepKind::CreateIndexConcurrent,
175                TransactionConstraint::OutsideTransaction,
176            ),
177            step(StepKind::CreateTable, TransactionConstraint::InTransaction),
178            step(
179                StepKind::CreateIndexConcurrent,
180                TransactionConstraint::OutsideTransaction,
181            ),
182        ]);
183        let ids: Vec<u32> = out.iter().map(|g| g.id).collect();
184        assert_eq!(ids, vec![1, 2, 3, 4]);
185    }
186}
pgevolve_core/plan/grouping.rs

pgevolve_core/plan/
grouping.rs
