es_entity/
tree_query.rs

1//! Dynamically assembles the single-statement tree query a nested repo's
2//! generated read fns execute at runtime: a tagged `UNION ALL` over one CTE
3//! per tree node, so a parent and all of its nested children (recursively)
4//! come back in one statement instead of one per level.
5
6use std::collections::HashMap;
7
8use crate::{db, events::GenericEvent};
9
10/// Static description of one node (repo) in a nested tree, used to build the query.
11#[derive(Debug, Clone)]
12pub struct TreeSpec {
13    pub table_name: &'static str,
14    pub events_table_name: &'static str,
15    pub parent_column: Option<&'static str>,
16    pub soft_delete: bool,
17    pub forgettable_table_name: Option<&'static str>,
18    pub event_context: bool,
19    pub children: Vec<TreeSpec>,
20}
21
22/// The root user query plus the bits needed to fold it into the tree query.
23pub struct TreeQuerySource<Id> {
24    pub user_sql: &'static str,
25    pub order_by_cols: &'static [&'static str],
26    pub n_user_args: usize,
27    pub decode: fn(&db::Row) -> Result<GenericEvent<Id>, sqlx::Error>,
28}
29
30pub fn decode_tag(row: &db::Row) -> Result<i32, sqlx::Error> {
31    sqlx::Row::try_get(row, "tag")
32}
33
34pub fn decode_tagged_row<Id>(row: &db::Row) -> Result<GenericEvent<Id>, sqlx::Error>
35where
36    Id: for<'r> sqlx::Decode<'r, db::Db> + sqlx::Type<db::Db>,
37{
38    use sqlx::Row;
39    Ok(GenericEvent {
40        entity_id: row.try_get("entity_id")?,
41        sequence: row.try_get("sequence")?,
42        event: row.try_get("event")?,
43        context: row.try_get("context")?,
44        recorded_at: row.try_get("recorded_at")?,
45        forgettable_payload: row.try_get("forgettable_payload")?,
46    })
47}
48
49pub fn partition_by_tag(rows: Vec<db::Row>) -> Result<HashMap<i32, Vec<db::Row>>, sqlx::Error> {
50    let mut by_tag: HashMap<i32, Vec<db::Row>> = HashMap::new();
51    for row in rows {
52        let tag = decode_tag(&row)?;
53        by_tag.entry(tag).or_default().push(row);
54    }
55    Ok(by_tag)
56}
57
58fn branch_sql(
59    tag: i32,
60    cte_name: &str,
61    node: &TreeSpec,
62    is_root: bool,
63    ctx_param_idx: usize,
64) -> String {
65    let context_expr = if is_root {
66        format!("CASE WHEN ${ctx_param_idx} THEN e.context ELSE NULL::jsonb END")
67    } else if node.event_context {
68        "e.context".to_string()
69    } else {
70        "NULL::jsonb".to_string()
71    };
72    let (payload_expr, forgettable_join) = match node.forgettable_table_name {
73        Some(tbl) => (
74            "p.payload".to_string(),
75            format!(" LEFT JOIN {tbl} p ON e.id = p.entity_id AND e.sequence = p.sequence"),
76        ),
77        None => ("NULL::jsonb".to_string(), String::new()),
78    };
79    let ord_expr = if is_root { "i.__ord" } else { "NULL::BIGINT" };
80    format!(
81        "SELECT {tag} AS tag, i.id AS entity_id, e.sequence, e.event, {context_expr} AS context, \
82         e.recorded_at, {payload_expr} AS forgettable_payload, {ord_expr} AS __ord \
83         FROM {cte_name} i JOIN {events_table} e ON i.id = e.id{forgettable_join}",
84        events_table = node.events_table_name,
85    )
86}
87
88pub fn build_tree_query(
89    user_sql: &str,
90    order_by_cols: &[&str],
91    spec: &TreeSpec,
92    include_deleted: bool,
93    ctx_param_idx: usize,
94) -> String {
95    let order_clause = if order_by_cols.is_empty() {
96        "id".to_string()
97    } else {
98        order_by_cols.join(", ")
99    };
100
101    let mut ctes = vec![
102        format!("__user AS ({user_sql})"),
103        format!(
104            "entities AS (SELECT *, ROW_NUMBER() OVER (ORDER BY {order_clause}) AS __ord FROM __user)"
105        ),
106    ];
107    let mut branches = vec![branch_sql(0, "entities", spec, true, ctx_param_idx)];
108    let mut cursor: i32 = 1;
109
110    walk_children(
111        spec,
112        "entities",
113        include_deleted,
114        ctx_param_idx,
115        &mut cursor,
116        &mut ctes,
117        &mut branches,
118    );
119
120    format!(
121        "WITH {} {} ORDER BY tag, __ord, entity_id, sequence",
122        ctes.join(", "),
123        branches.join(" UNION ALL "),
124    )
125}
126
127#[allow(clippy::too_many_arguments)]
128fn walk_children(
129    node: &TreeSpec,
130    parent_cte: &str,
131    include_deleted: bool,
132    ctx_param_idx: usize,
133    cursor: &mut i32,
134    ctes: &mut Vec<String>,
135    branches: &mut Vec<String>,
136) {
137    for child in &node.children {
138        let tag = *cursor;
139        *cursor += 1;
140        let cte_name = format!("n{tag}");
141        let parent_col = child
142            .parent_column
143            .expect("non-root tree node must declare a parent column");
144        let deleted_cond = if child.soft_delete && !include_deleted {
145            " AND deleted = FALSE"
146        } else {
147            ""
148        };
149        ctes.push(format!(
150            "{cte_name} AS (SELECT id FROM {} WHERE {parent_col} IN (SELECT id FROM {parent_cte}){deleted_cond})",
151            child.table_name
152        ));
153        branches.push(branch_sql(tag, &cte_name, child, false, ctx_param_idx));
154        walk_children(
155            child,
156            &cte_name,
157            include_deleted,
158            ctx_param_idx,
159            cursor,
160            ctes,
161            branches,
162        );
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    fn leaf(table: &'static str, events: &'static str, parent_col: &'static str) -> TreeSpec {
171        TreeSpec {
172            table_name: table,
173            events_table_name: events,
174            parent_column: Some(parent_col),
175            soft_delete: false,
176            forgettable_table_name: None,
177            event_context: false,
178            children: Vec::new(),
179        }
180    }
181
182    fn root(table: &'static str, events: &'static str, children: Vec<TreeSpec>) -> TreeSpec {
183        TreeSpec {
184            table_name: table,
185            events_table_name: events,
186            parent_column: None,
187            soft_delete: false,
188            forgettable_table_name: None,
189            event_context: false,
190            children,
191        }
192    }
193
194    #[test]
195    fn leaf_only_root() {
196        let spec = root("subscriptions", "subscription_events", Vec::new());
197        let sql = build_tree_query(
198            "SELECT id FROM subscriptions WHERE id = $1",
199            &[],
200            &spec,
201            false,
202            2,
203        );
204        assert_eq!(
205            sql,
206            "WITH __user AS (SELECT id FROM subscriptions WHERE id = $1), \
207             entities AS (SELECT *, ROW_NUMBER() OVER (ORDER BY id) AS __ord FROM __user) \
208             SELECT 0 AS tag, i.id AS entity_id, e.sequence, e.event, \
209             CASE WHEN $2 THEN e.context ELSE NULL::jsonb END AS context, e.recorded_at, \
210             NULL::jsonb AS forgettable_payload, i.__ord AS __ord \
211             FROM entities i JOIN subscription_events e ON i.id = e.id \
212             ORDER BY tag, __ord, entity_id, sequence"
213        );
214    }
215
216    #[test]
217    fn one_child() {
218        let child = TreeSpec {
219            soft_delete: true,
220            ..leaf(
221                "billing_periods",
222                "billing_period_events",
223                "subscription_id",
224            )
225        };
226        let spec = root("subscriptions", "subscription_events", vec![child]);
227        let sql = build_tree_query(
228            "SELECT id FROM subscriptions WHERE id = $1",
229            &[],
230            &spec,
231            false,
232            2,
233        );
234        assert!(sql.contains(
235            "n1 AS (SELECT id FROM billing_periods WHERE subscription_id IN (SELECT id FROM entities) AND deleted = FALSE)"
236        ));
237        assert!(sql.contains(
238            "SELECT 1 AS tag, i.id AS entity_id, e.sequence, e.event, NULL::jsonb AS context, \
239             e.recorded_at, NULL::jsonb AS forgettable_payload, NULL::BIGINT AS __ord \
240             FROM n1 i JOIN billing_period_events e ON i.id = e.id"
241        ));
242    }
243
244    #[test]
245    fn two_children_first_has_grandchild() {
246        let grandchild = leaf("line_items", "line_item_events", "billing_period_id");
247        let child_a = TreeSpec {
248            children: vec![grandchild],
249            ..leaf(
250                "billing_periods",
251                "billing_period_events",
252                "subscription_id",
253            )
254        };
255        let child_b = leaf("invoices", "invoice_events", "subscription_id");
256        let spec = root(
257            "subscriptions",
258            "subscription_events",
259            vec![child_a, child_b],
260        );
261        let sql = build_tree_query(
262            "SELECT id FROM subscriptions WHERE id = $1",
263            &[],
264            &spec,
265            false,
266            2,
267        );
268        assert!(sql.contains("n1 AS (SELECT id FROM billing_periods WHERE subscription_id IN (SELECT id FROM entities)"));
269        assert!(sql.contains(
270            "n2 AS (SELECT id FROM line_items WHERE billing_period_id IN (SELECT id FROM n1)"
271        ));
272        assert!(sql.contains(
273            "n3 AS (SELECT id FROM invoices WHERE subscription_id IN (SELECT id FROM entities)"
274        ));
275        assert!(sql.contains("SELECT 1 AS tag"));
276        assert!(sql.contains("SELECT 2 AS tag"));
277        assert!(sql.contains("SELECT 3 AS tag"));
278    }
279
280    #[test]
281    fn soft_delete_and_forgettable_and_inlined_context() {
282        let child = TreeSpec {
283            soft_delete: true,
284            forgettable_table_name: Some("order_items_forgettable_payloads"),
285            event_context: true,
286            ..leaf("order_items", "order_item_events", "order_id")
287        };
288        let spec = root("orders", "order_events", vec![child]);
289        let sql = build_tree_query("SELECT id FROM orders WHERE id = $1", &[], &spec, false, 2);
290        assert!(sql.contains("AND deleted = FALSE)"));
291        assert!(sql.contains(
292            "LEFT JOIN order_items_forgettable_payloads p ON e.id = p.entity_id AND e.sequence = p.sequence"
293        ));
294        assert!(sql.contains("p.payload AS forgettable_payload"));
295        assert!(sql.contains(
296            "SELECT 1 AS tag, i.id AS entity_id, e.sequence, e.event, e.context AS context"
297        ));
298    }
299
300    #[test]
301    fn include_deleted_drops_deleted_condition_transitively() {
302        let grandchild = TreeSpec {
303            soft_delete: true,
304            ..leaf("line_items", "line_item_events", "billing_period_id")
305        };
306        let child = TreeSpec {
307            soft_delete: true,
308            children: vec![grandchild],
309            ..leaf(
310                "billing_periods",
311                "billing_period_events",
312                "subscription_id",
313            )
314        };
315        let spec = root("subscriptions", "subscription_events", vec![child]);
316        let sql = build_tree_query(
317            "SELECT id FROM subscriptions WHERE id = $1",
318            &[],
319            &spec,
320            true,
321            2,
322        );
323        assert!(!sql.contains("deleted = FALSE"));
324    }
325
326    #[test]
327    fn empty_order_by_defaults_to_id() {
328        let spec = root("subscriptions", "subscription_events", Vec::new());
329        let sql = build_tree_query(
330            "SELECT id FROM subscriptions WHERE id = $1",
331            &[],
332            &spec,
333            false,
334            2,
335        );
336        assert!(sql.contains("ROW_NUMBER() OVER (ORDER BY id) AS __ord"));
337    }
338
339    #[test]
340    fn explicit_order_by_is_unprefixed_in_the_entities_cte() {
341        let spec = root("entities", "entity_events", Vec::new());
342        let sql = build_tree_query(
343            "SELECT name, id FROM entities ORDER BY name, id LIMIT $1",
344            &["name", "id"],
345            &spec,
346            false,
347            2,
348        );
349        assert!(sql.contains("ROW_NUMBER() OVER (ORDER BY name, id) AS __ord"));
350    }
351}
es_entity/tree_query.rs

es_entity/
tree_query.rs
