Skip to main content

qail_core/optimizer/
normalized_mutation.rs

1use crate::ast::{
2    Action, Cage, CageKind, Condition, Expr, LogicalOp, OnConflict, OverridingKind, Qail, Value,
3};
4use std::collections::HashSet;
5
6/// Canonical condition clause used by mutation normalization.
7#[derive(Debug, Clone, PartialEq)]
8pub struct MutationClause {
9    pub logical_op: LogicalOp,
10    pub conditions: Vec<Condition>,
11}
12
13/// Canonical representation for the rewrite-safe subset of mutation queries.
14#[derive(Debug, Clone, PartialEq)]
15pub struct NormalizedMutation {
16    pub action: Action,
17    pub table: String,
18    pub columns: Vec<Expr>,
19    pub payload: Vec<MutationClause>,
20    pub filters: Vec<MutationClause>,
21    pub returning: Option<Vec<Expr>>,
22    pub on_conflict: Option<OnConflict>,
23    pub source_query: Option<Box<Qail>>,
24    pub default_values: bool,
25    pub overriding: Option<OverridingKind>,
26    pub from_tables: Vec<String>,
27    pub using_tables: Vec<String>,
28    pub only_table: bool,
29}
30
31/// Errors returned when a query cannot be normalized into the supported mutation subset.
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum NormalizeMutationError {
34    UnsupportedAction(Action),
35    UnsupportedFeature(&'static str),
36    InvalidShape(&'static str),
37}
38
39impl std::fmt::Display for NormalizeMutationError {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        match self {
42            Self::UnsupportedAction(action) => {
43                write!(
44                    f,
45                    "mutation normalization only supports ADD/SET/DEL, got {}",
46                    action
47                )
48            }
49            Self::UnsupportedFeature(feature) => {
50                write!(f, "unsupported mutation feature: {}", feature)
51            }
52            Self::InvalidShape(shape) => write!(f, "invalid mutation shape: {}", shape),
53        }
54    }
55}
56
57impl std::error::Error for NormalizeMutationError {}
58
59/// Normalize a QAIL mutation (`ADD`, `SET`, `DEL`) into a canonical representation.
60pub fn normalize_mutation(qail: &Qail) -> Result<NormalizedMutation, NormalizeMutationError> {
61    NormalizedMutation::try_from(qail)
62}
63
64impl TryFrom<&Qail> for NormalizedMutation {
65    type Error = NormalizeMutationError;
66
67    fn try_from(qail: &Qail) -> Result<Self, Self::Error> {
68        if !matches!(qail.action, Action::Add | Action::Set | Action::Del) {
69            return Err(NormalizeMutationError::UnsupportedAction(qail.action));
70        }
71
72        reject_unsupported_mutation_features(qail)?;
73
74        let mut payload = Vec::new();
75        let mut filters = Vec::new();
76        for cage in &qail.cages {
77            match &cage.kind {
78                CageKind::Payload => payload.push(MutationClause {
79                    logical_op: cage.logical_op,
80                    conditions: cage.conditions.clone(),
81                }),
82                CageKind::Filter => filters.push(MutationClause {
83                    logical_op: cage.logical_op,
84                    conditions: cage.conditions.clone(),
85                }),
86                CageKind::Sort(_) => {
87                    return Err(NormalizeMutationError::UnsupportedFeature("ORDER BY cages"));
88                }
89                CageKind::Limit(_) => {
90                    return Err(NormalizeMutationError::UnsupportedFeature("LIMIT cages"));
91                }
92                CageKind::Offset(_) => {
93                    return Err(NormalizeMutationError::UnsupportedFeature("OFFSET cages"));
94                }
95                CageKind::Sample(_) => {
96                    return Err(NormalizeMutationError::UnsupportedFeature("sample cages"));
97                }
98                CageKind::Qualify => {
99                    return Err(NormalizeMutationError::UnsupportedFeature("QUALIFY cages"));
100                }
101                CageKind::Partition => {
102                    return Err(NormalizeMutationError::UnsupportedFeature("GROUP BY cages"));
103                }
104            }
105        }
106
107        match qail.action {
108            Action::Add => {
109                if !qail.from_tables.is_empty() || !qail.using_tables.is_empty() || qail.only_table
110                {
111                    return Err(NormalizeMutationError::UnsupportedFeature(
112                        "INSERT with FROM/USING/ONLY",
113                    ));
114                }
115
116                if !filters.is_empty() {
117                    return Err(NormalizeMutationError::UnsupportedFeature(
118                        "INSERT filter cages",
119                    ));
120                }
121
122                let inserts_from_non_payload =
123                    qail.default_values || qail.source_query.is_some() || qail.columns.is_empty();
124
125                if inserts_from_non_payload {
126                    if !payload.is_empty() {
127                        return Err(NormalizeMutationError::InvalidShape(
128                            "payload cages with DEFAULT VALUES or INSERT ... SELECT",
129                        ));
130                    }
131                } else if payload.len() != 1 {
132                    return Err(NormalizeMutationError::InvalidShape(
133                        "INSERT requires exactly one payload cage",
134                    ));
135                }
136            }
137            Action::Set => {
138                if !qail.using_tables.is_empty() {
139                    return Err(NormalizeMutationError::UnsupportedFeature("UPDATE USING"));
140                }
141                if qail.on_conflict.is_some() {
142                    return Err(NormalizeMutationError::UnsupportedFeature(
143                        "UPDATE ON CONFLICT",
144                    ));
145                }
146                if qail.source_query.is_some() {
147                    return Err(NormalizeMutationError::UnsupportedFeature(
148                        "UPDATE source query",
149                    ));
150                }
151                if qail.default_values {
152                    return Err(NormalizeMutationError::UnsupportedFeature(
153                        "UPDATE DEFAULT VALUES",
154                    ));
155                }
156                if qail.overriding.is_some() {
157                    return Err(NormalizeMutationError::UnsupportedFeature(
158                        "UPDATE OVERRIDING",
159                    ));
160                }
161            }
162            Action::Del => {
163                if !qail.from_tables.is_empty() {
164                    return Err(NormalizeMutationError::UnsupportedFeature("DELETE FROM"));
165                }
166                if !payload.is_empty() {
167                    return Err(NormalizeMutationError::UnsupportedFeature(
168                        "DELETE payload cages",
169                    ));
170                }
171                if qail.on_conflict.is_some() {
172                    return Err(NormalizeMutationError::UnsupportedFeature(
173                        "DELETE ON CONFLICT",
174                    ));
175                }
176                if qail.source_query.is_some() {
177                    return Err(NormalizeMutationError::UnsupportedFeature(
178                        "DELETE source query",
179                    ));
180                }
181                if qail.default_values {
182                    return Err(NormalizeMutationError::UnsupportedFeature(
183                        "DELETE DEFAULT VALUES",
184                    ));
185                }
186                if qail.overriding.is_some() {
187                    return Err(NormalizeMutationError::UnsupportedFeature(
188                        "DELETE OVERRIDING",
189                    ));
190                }
191            }
192            _ => unreachable!("unsupported action already rejected"),
193        }
194
195        Ok(Self {
196            action: qail.action,
197            table: qail.table.clone(),
198            columns: qail.columns.clone(),
199            payload,
200            filters,
201            returning: qail.returning.clone(),
202            on_conflict: qail.on_conflict.clone(),
203            source_query: qail.source_query.clone(),
204            default_values: qail.default_values,
205            overriding: qail.overriding,
206            from_tables: qail.from_tables.clone(),
207            using_tables: qail.using_tables.clone(),
208            only_table: qail.only_table,
209        })
210    }
211}
212
213impl NormalizedMutation {
214    /// Apply deterministic cleanup on supported mutation shapes.
215    ///
216    /// Cleanup rules:
217    /// - merge multiple `SET` payload cages into one (preserving assignment order)
218    /// - merge all `AND` filter clauses into one clause
219    /// - merge all `OR` filter clauses into one clause
220    /// - remove empty filter clauses
221    /// - sort and dedupe filter conditions
222    pub fn cleaned(&self) -> Self {
223        let mut cleaned = self.clone();
224
225        if cleaned.action == Action::Set {
226            let mut merged_payload = Vec::new();
227            for clause in &cleaned.payload {
228                merged_payload.extend(clause.conditions.clone());
229            }
230            cleaned.payload = if merged_payload.is_empty() {
231                Vec::new()
232            } else {
233                vec![MutationClause {
234                    logical_op: LogicalOp::And,
235                    conditions: merged_payload,
236                }]
237            };
238        }
239
240        let mut and_conditions = Vec::new();
241        let mut or_conditions = Vec::new();
242        for filter in &cleaned.filters {
243            match filter.logical_op {
244                LogicalOp::And => and_conditions.extend(filter.conditions.clone()),
245                LogicalOp::Or => or_conditions.extend(filter.conditions.clone()),
246            }
247        }
248
249        and_conditions = dedupe_conditions_sorted(and_conditions);
250        or_conditions = dedupe_conditions_sorted(or_conditions);
251
252        cleaned.filters.clear();
253        if !and_conditions.is_empty() {
254            cleaned.filters.push(MutationClause {
255                logical_op: LogicalOp::And,
256                conditions: and_conditions,
257            });
258        }
259        if !or_conditions.is_empty() {
260            cleaned.filters.push(MutationClause {
261                logical_op: LogicalOp::Or,
262                conditions: or_conditions,
263            });
264        }
265
266        cleaned
267    }
268
269    /// Return a canonicalized clone for structural comparison.
270    pub fn canonicalized(&self) -> Self {
271        self.cleaned()
272    }
273
274    /// Compare two normalized mutations under the supported-subset semantics.
275    pub fn equivalent_shape(&self, other: &Self) -> bool {
276        self.canonicalized() == other.canonicalized()
277    }
278
279    /// Lower the normalized form back into canonical QAIL.
280    pub fn to_qail(&self) -> Qail {
281        let mut qail = Qail {
282            action: self.action,
283            table: self.table.clone(),
284            columns: self.columns.clone(),
285            returning: self.returning.clone(),
286            on_conflict: self.on_conflict.clone(),
287            source_query: self.source_query.clone(),
288            default_values: self.default_values,
289            overriding: self.overriding,
290            from_tables: self.from_tables.clone(),
291            using_tables: self.using_tables.clone(),
292            only_table: self.only_table,
293            ..Default::default()
294        };
295
296        for payload in &self.payload {
297            qail.cages.push(Cage {
298                kind: CageKind::Payload,
299                conditions: payload.conditions.clone(),
300                logical_op: payload.logical_op,
301            });
302        }
303
304        for filter in &self.filters {
305            qail.cages.push(Cage {
306                kind: CageKind::Filter,
307                conditions: filter.conditions.clone(),
308                logical_op: filter.logical_op,
309            });
310        }
311
312        qail
313    }
314}
315
316fn reject_unsupported_mutation_features(qail: &Qail) -> Result<(), NormalizeMutationError> {
317    if !qail.joins.is_empty() {
318        return Err(NormalizeMutationError::UnsupportedFeature("joins"));
319    }
320    if qail.distinct {
321        return Err(NormalizeMutationError::UnsupportedFeature("DISTINCT"));
322    }
323    if qail.index_def.is_some() {
324        return Err(NormalizeMutationError::UnsupportedFeature(
325            "index definitions",
326        ));
327    }
328    if !qail.table_constraints.is_empty() {
329        return Err(NormalizeMutationError::UnsupportedFeature(
330            "table constraints",
331        ));
332    }
333    if !qail.set_ops.is_empty() {
334        return Err(NormalizeMutationError::UnsupportedFeature("set operations"));
335    }
336    if !qail.having.is_empty() {
337        return Err(NormalizeMutationError::UnsupportedFeature("HAVING"));
338    }
339    if !qail.group_by_mode.is_simple() {
340        return Err(NormalizeMutationError::UnsupportedFeature("GROUP BY mode"));
341    }
342    if !qail.ctes.is_empty() {
343        return Err(NormalizeMutationError::UnsupportedFeature("CTEs"));
344    }
345    if !qail.distinct_on.is_empty() {
346        return Err(NormalizeMutationError::UnsupportedFeature("DISTINCT ON"));
347    }
348    if qail.channel.is_some() || qail.payload.is_some() {
349        return Err(NormalizeMutationError::UnsupportedFeature(
350            "LISTEN/NOTIFY metadata",
351        ));
352    }
353    if qail.savepoint_name.is_some() {
354        return Err(NormalizeMutationError::UnsupportedFeature(
355            "savepoint metadata",
356        ));
357    }
358    if qail.lock_mode.is_some() || qail.skip_locked {
359        return Err(NormalizeMutationError::UnsupportedFeature("row locks"));
360    }
361    if qail.fetch.is_some() {
362        return Err(NormalizeMutationError::UnsupportedFeature("FETCH FIRST"));
363    }
364    if qail.sample.is_some() {
365        return Err(NormalizeMutationError::UnsupportedFeature("TABLESAMPLE"));
366    }
367    if qail.vector.is_some()
368        || qail.score_threshold.is_some()
369        || qail.vector_name.is_some()
370        || qail.with_vector
371        || qail.vector_size.is_some()
372        || qail.distance.is_some()
373        || qail.on_disk.is_some()
374    {
375        return Err(NormalizeMutationError::UnsupportedFeature(
376            "vector search fields",
377        ));
378    }
379    if qail.function_def.is_some() || qail.trigger_def.is_some() {
380        return Err(NormalizeMutationError::UnsupportedFeature(
381            "procedural objects",
382        ));
383    }
384
385    Ok(())
386}
387
388fn condition_signature(condition: &Condition) -> String {
389    format!(
390        "{}|{}|{}|{}",
391        expr_signature(&condition.left),
392        condition.op.sql_symbol(),
393        value_signature(&condition.value),
394        condition.is_array_unnest
395    )
396}
397
398fn expr_signature(expr: &Expr) -> String {
399    format!("{}", expr)
400}
401
402fn value_signature(value: &Value) -> String {
403    format!("{}", value)
404}
405
406fn dedupe_conditions_sorted(mut conditions: Vec<Condition>) -> Vec<Condition> {
407    conditions.sort_by_key(condition_signature);
408
409    let mut seen = HashSet::new();
410    let mut deduped = Vec::with_capacity(conditions.len());
411    for condition in conditions {
412        let signature = condition_signature(&condition);
413        if seen.insert(signature) {
414            deduped.push(condition);
415        }
416    }
417    deduped
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423    use crate::ast::{Cage, ConflictAction, Operator};
424
425    fn eq(col: &str, value: Value) -> Condition {
426        Condition {
427            left: Expr::Named(col.to_string()),
428            op: Operator::Eq,
429            value,
430            is_array_unnest: false,
431        }
432    }
433
434    #[test]
435    fn normalize_supported_insert_shape() {
436        let qail = Qail {
437            action: Action::Add,
438            table: "users".to_string(),
439            columns: vec![
440                Expr::Named("id".to_string()),
441                Expr::Named("email".to_string()),
442            ],
443            cages: vec![Cage {
444                kind: CageKind::Payload,
445                logical_op: LogicalOp::And,
446                conditions: vec![
447                    eq("id", Value::Int(1)),
448                    eq("email", Value::String("a@b.com".to_string())),
449                ],
450            }],
451            on_conflict: Some(OnConflict {
452                columns: vec!["id".to_string()],
453                action: ConflictAction::DoNothing,
454            }),
455            returning: Some(vec![Expr::Star]),
456            ..Default::default()
457        };
458
459        let normalized = normalize_mutation(&qail).expect("insert should normalize");
460        assert_eq!(normalized.action, Action::Add);
461        assert_eq!(normalized.table, "users");
462        assert_eq!(normalized.payload.len(), 1);
463        assert!(normalized.filters.is_empty());
464        assert!(normalized.on_conflict.is_some());
465    }
466
467    #[test]
468    fn normalize_rejects_insert_with_multiple_payload_cages() {
469        let qail = Qail {
470            action: Action::Add,
471            table: "users".to_string(),
472            columns: vec![Expr::Named("id".to_string())],
473            cages: vec![
474                Cage {
475                    kind: CageKind::Payload,
476                    logical_op: LogicalOp::And,
477                    conditions: vec![eq("id", Value::Int(1))],
478                },
479                Cage {
480                    kind: CageKind::Payload,
481                    logical_op: LogicalOp::And,
482                    conditions: vec![eq("id", Value::Int(2))],
483                },
484            ],
485            ..Default::default()
486        };
487
488        let err = normalize_mutation(&qail).expect_err("multiple payload cages must be rejected");
489        assert_eq!(
490            err,
491            NormalizeMutationError::InvalidShape("INSERT requires exactly one payload cage")
492        );
493    }
494
495    #[test]
496    fn normalize_supported_update_shape() {
497        let qail = Qail {
498            action: Action::Set,
499            table: "users".to_string(),
500            cages: vec![
501                Cage {
502                    kind: CageKind::Payload,
503                    logical_op: LogicalOp::And,
504                    conditions: vec![eq("name", Value::String("Alice".to_string()))],
505                },
506                Cage {
507                    kind: CageKind::Filter,
508                    logical_op: LogicalOp::And,
509                    conditions: vec![eq("id", Value::Int(7))],
510                },
511            ],
512            from_tables: vec!["teams".to_string()],
513            ..Default::default()
514        };
515
516        let normalized = normalize_mutation(&qail).expect("update should normalize");
517        assert_eq!(normalized.action, Action::Set);
518        assert_eq!(normalized.from_tables, vec!["teams".to_string()]);
519        assert_eq!(normalized.payload.len(), 1);
520        assert_eq!(normalized.filters.len(), 1);
521    }
522
523    #[test]
524    fn normalize_supported_delete_shape() {
525        let qail = Qail {
526            action: Action::Del,
527            table: "users".to_string(),
528            cages: vec![Cage {
529                kind: CageKind::Filter,
530                logical_op: LogicalOp::And,
531                conditions: vec![eq("id", Value::Int(9))],
532            }],
533            using_tables: vec!["teams".to_string()],
534            only_table: true,
535            ..Default::default()
536        };
537
538        let normalized = normalize_mutation(&qail).expect("delete should normalize");
539        assert_eq!(normalized.action, Action::Del);
540        assert_eq!(normalized.using_tables, vec!["teams".to_string()]);
541        assert!(normalized.payload.is_empty());
542        assert!(normalized.only_table);
543    }
544
545    #[test]
546    fn cleanup_merges_update_payload_and_filter_clauses() {
547        let qail = Qail {
548            action: Action::Set,
549            table: "users".to_string(),
550            cages: vec![
551                Cage {
552                    kind: CageKind::Payload,
553                    logical_op: LogicalOp::And,
554                    conditions: vec![eq("name", Value::String("Alice".to_string()))],
555                },
556                Cage {
557                    kind: CageKind::Payload,
558                    logical_op: LogicalOp::And,
559                    conditions: vec![eq("active", Value::Bool(true))],
560                },
561                Cage {
562                    kind: CageKind::Filter,
563                    logical_op: LogicalOp::And,
564                    conditions: vec![eq("id", Value::Int(1)), eq("id", Value::Int(1))],
565                },
566                Cage {
567                    kind: CageKind::Filter,
568                    logical_op: LogicalOp::Or,
569                    conditions: vec![eq("role", Value::String("admin".to_string()))],
570                },
571            ],
572            ..Default::default()
573        };
574
575        let normalized = normalize_mutation(&qail).expect("update should normalize");
576        let cleaned = normalized.cleaned();
577
578        assert_eq!(cleaned.payload.len(), 1);
579        assert_eq!(cleaned.payload[0].conditions.len(), 2);
580        assert_eq!(cleaned.filters.len(), 2);
581        assert_eq!(cleaned.filters[0].logical_op, LogicalOp::And);
582        assert_eq!(cleaned.filters[0].conditions.len(), 1);
583        assert_eq!(cleaned.filters[1].logical_op, LogicalOp::Or);
584    }
585
586    #[test]
587    fn normalized_mutation_roundtrips_to_canonical_qail() {
588        let qail = Qail {
589            action: Action::Set,
590            table: "users".to_string(),
591            cages: vec![
592                Cage {
593                    kind: CageKind::Payload,
594                    logical_op: LogicalOp::And,
595                    conditions: vec![eq("email", Value::String("x@y.com".to_string()))],
596                },
597                Cage {
598                    kind: CageKind::Filter,
599                    logical_op: LogicalOp::And,
600                    conditions: vec![eq("id", Value::Int(42))],
601                },
602            ],
603            ..Default::default()
604        };
605
606        let normalized = normalize_mutation(&qail).expect("update should normalize");
607        let roundtrip = normalized.to_qail();
608        let roundtrip_normalized =
609            normalize_mutation(&roundtrip).expect("roundtrip should normalize");
610
611        assert!(normalized.equivalent_shape(&roundtrip_normalized));
612    }
613
614    #[test]
615    fn cleanup_is_idempotent() {
616        let qail = Qail {
617            action: Action::Del,
618            table: "users".to_string(),
619            cages: vec![
620                Cage {
621                    kind: CageKind::Filter,
622                    logical_op: LogicalOp::And,
623                    conditions: vec![eq("id", Value::Int(1)), eq("id", Value::Int(1))],
624                },
625                Cage {
626                    kind: CageKind::Filter,
627                    logical_op: LogicalOp::And,
628                    conditions: vec![eq("active", Value::Bool(true))],
629                },
630            ],
631            ..Default::default()
632        };
633
634        let normalized = normalize_mutation(&qail).expect("delete should normalize");
635        let cleaned = normalized.cleaned();
636        let cleaned_twice = cleaned.cleaned();
637        assert_eq!(cleaned, cleaned_twice);
638    }
639
640    #[test]
641    fn equivalent_shape_ignores_filter_condition_order() {
642        let left = Qail {
643            action: Action::Del,
644            table: "users".to_string(),
645            cages: vec![Cage {
646                kind: CageKind::Filter,
647                logical_op: LogicalOp::And,
648                conditions: vec![
649                    eq("active", Value::Bool(true)),
650                    eq("tenant_id", Value::String("t1".to_string())),
651                ],
652            }],
653            ..Default::default()
654        };
655        let right = Qail {
656            action: Action::Del,
657            table: "users".to_string(),
658            cages: vec![Cage {
659                kind: CageKind::Filter,
660                logical_op: LogicalOp::And,
661                conditions: vec![
662                    eq("tenant_id", Value::String("t1".to_string())),
663                    eq("active", Value::Bool(true)),
664                ],
665            }],
666            ..Default::default()
667        };
668
669        let left = normalize_mutation(&left).expect("left mutation should normalize");
670        let right = normalize_mutation(&right).expect("right mutation should normalize");
671
672        assert!(left.equivalent_shape(&right));
673    }
674}