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            _ => return Err(NormalizeMutationError::UnsupportedAction(qail.action)),
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                where_conditions: Vec::new(),
455            }),
456            returning: Some(vec![Expr::Star]),
457            ..Default::default()
458        };
459
460        let normalized = normalize_mutation(&qail).expect("insert should normalize");
461        assert_eq!(normalized.action, Action::Add);
462        assert_eq!(normalized.table, "users");
463        assert_eq!(normalized.payload.len(), 1);
464        assert!(normalized.filters.is_empty());
465        assert!(normalized.on_conflict.is_some());
466    }
467
468    #[test]
469    fn normalize_rejects_insert_with_multiple_payload_cages() {
470        let qail = Qail {
471            action: Action::Add,
472            table: "users".to_string(),
473            columns: vec![Expr::Named("id".to_string())],
474            cages: vec![
475                Cage {
476                    kind: CageKind::Payload,
477                    logical_op: LogicalOp::And,
478                    conditions: vec![eq("id", Value::Int(1))],
479                },
480                Cage {
481                    kind: CageKind::Payload,
482                    logical_op: LogicalOp::And,
483                    conditions: vec![eq("id", Value::Int(2))],
484                },
485            ],
486            ..Default::default()
487        };
488
489        let err = normalize_mutation(&qail).expect_err("multiple payload cages must be rejected");
490        assert_eq!(
491            err,
492            NormalizeMutationError::InvalidShape("INSERT requires exactly one payload cage")
493        );
494    }
495
496    #[test]
497    fn normalize_supported_update_shape() {
498        let qail = Qail {
499            action: Action::Set,
500            table: "users".to_string(),
501            cages: vec![
502                Cage {
503                    kind: CageKind::Payload,
504                    logical_op: LogicalOp::And,
505                    conditions: vec![eq("name", Value::String("Alice".to_string()))],
506                },
507                Cage {
508                    kind: CageKind::Filter,
509                    logical_op: LogicalOp::And,
510                    conditions: vec![eq("id", Value::Int(7))],
511                },
512            ],
513            from_tables: vec!["teams".to_string()],
514            ..Default::default()
515        };
516
517        let normalized = normalize_mutation(&qail).expect("update should normalize");
518        assert_eq!(normalized.action, Action::Set);
519        assert_eq!(normalized.from_tables, vec!["teams".to_string()]);
520        assert_eq!(normalized.payload.len(), 1);
521        assert_eq!(normalized.filters.len(), 1);
522    }
523
524    #[test]
525    fn normalize_supported_delete_shape() {
526        let qail = Qail {
527            action: Action::Del,
528            table: "users".to_string(),
529            cages: vec![Cage {
530                kind: CageKind::Filter,
531                logical_op: LogicalOp::And,
532                conditions: vec![eq("id", Value::Int(9))],
533            }],
534            using_tables: vec!["teams".to_string()],
535            only_table: true,
536            ..Default::default()
537        };
538
539        let normalized = normalize_mutation(&qail).expect("delete should normalize");
540        assert_eq!(normalized.action, Action::Del);
541        assert_eq!(normalized.using_tables, vec!["teams".to_string()]);
542        assert!(normalized.payload.is_empty());
543        assert!(normalized.only_table);
544    }
545
546    #[test]
547    fn cleanup_merges_update_payload_and_filter_clauses() {
548        let qail = Qail {
549            action: Action::Set,
550            table: "users".to_string(),
551            cages: vec![
552                Cage {
553                    kind: CageKind::Payload,
554                    logical_op: LogicalOp::And,
555                    conditions: vec![eq("name", Value::String("Alice".to_string()))],
556                },
557                Cage {
558                    kind: CageKind::Payload,
559                    logical_op: LogicalOp::And,
560                    conditions: vec![eq("active", Value::Bool(true))],
561                },
562                Cage {
563                    kind: CageKind::Filter,
564                    logical_op: LogicalOp::And,
565                    conditions: vec![eq("id", Value::Int(1)), eq("id", Value::Int(1))],
566                },
567                Cage {
568                    kind: CageKind::Filter,
569                    logical_op: LogicalOp::Or,
570                    conditions: vec![eq("role", Value::String("admin".to_string()))],
571                },
572            ],
573            ..Default::default()
574        };
575
576        let normalized = normalize_mutation(&qail).expect("update should normalize");
577        let cleaned = normalized.cleaned();
578
579        assert_eq!(cleaned.payload.len(), 1);
580        assert_eq!(cleaned.payload[0].conditions.len(), 2);
581        assert_eq!(cleaned.filters.len(), 2);
582        assert_eq!(cleaned.filters[0].logical_op, LogicalOp::And);
583        assert_eq!(cleaned.filters[0].conditions.len(), 1);
584        assert_eq!(cleaned.filters[1].logical_op, LogicalOp::Or);
585    }
586
587    #[test]
588    fn normalized_mutation_roundtrips_to_canonical_qail() {
589        let qail = Qail {
590            action: Action::Set,
591            table: "users".to_string(),
592            cages: vec![
593                Cage {
594                    kind: CageKind::Payload,
595                    logical_op: LogicalOp::And,
596                    conditions: vec![eq("email", Value::String("x@y.com".to_string()))],
597                },
598                Cage {
599                    kind: CageKind::Filter,
600                    logical_op: LogicalOp::And,
601                    conditions: vec![eq("id", Value::Int(42))],
602                },
603            ],
604            ..Default::default()
605        };
606
607        let normalized = normalize_mutation(&qail).expect("update should normalize");
608        let roundtrip = normalized.to_qail();
609        let roundtrip_normalized =
610            normalize_mutation(&roundtrip).expect("roundtrip should normalize");
611
612        assert!(normalized.equivalent_shape(&roundtrip_normalized));
613    }
614
615    #[test]
616    fn cleanup_is_idempotent() {
617        let qail = Qail {
618            action: Action::Del,
619            table: "users".to_string(),
620            cages: vec![
621                Cage {
622                    kind: CageKind::Filter,
623                    logical_op: LogicalOp::And,
624                    conditions: vec![eq("id", Value::Int(1)), eq("id", Value::Int(1))],
625                },
626                Cage {
627                    kind: CageKind::Filter,
628                    logical_op: LogicalOp::And,
629                    conditions: vec![eq("active", Value::Bool(true))],
630                },
631            ],
632            ..Default::default()
633        };
634
635        let normalized = normalize_mutation(&qail).expect("delete should normalize");
636        let cleaned = normalized.cleaned();
637        let cleaned_twice = cleaned.cleaned();
638        assert_eq!(cleaned, cleaned_twice);
639    }
640
641    #[test]
642    fn equivalent_shape_ignores_filter_condition_order() {
643        let left = Qail {
644            action: Action::Del,
645            table: "users".to_string(),
646            cages: vec![Cage {
647                kind: CageKind::Filter,
648                logical_op: LogicalOp::And,
649                conditions: vec![
650                    eq("active", Value::Bool(true)),
651                    eq("tenant_id", Value::String("t1".to_string())),
652                ],
653            }],
654            ..Default::default()
655        };
656        let right = Qail {
657            action: Action::Del,
658            table: "users".to_string(),
659            cages: vec![Cage {
660                kind: CageKind::Filter,
661                logical_op: LogicalOp::And,
662                conditions: vec![
663                    eq("tenant_id", Value::String("t1".to_string())),
664                    eq("active", Value::Bool(true)),
665                ],
666            }],
667            ..Default::default()
668        };
669
670        let left = normalize_mutation(&left).expect("left mutation should normalize");
671        let right = normalize_mutation(&right).expect("right mutation should normalize");
672
673        assert!(left.equivalent_shape(&right));
674    }
675}