Skip to main content

uqa_sql/semantics/
returning.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! RETURNING target schemas, alias visibility, projection expansion, and static analysis.
8use super::returning_expression_schema;
9use crate::{
10    ast::{ColumnDef, ReturningAliases, Statement},
11    binding::snapshot::BindingSnapshot,
12    plan::{AggregateClassifier, ProjectionPlan},
13    routines::RoutineResolution,
14    ResultRow as Document, RowSchema, SQLError, SQLParam,
15};
16use std::collections::BTreeSet;
17use uqa_core::{DocId, Value};
18
19pub trait ReturningCatalog {
20    fn try_describe_table_row_type(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
21    fn try_table_columns(&self, table: &str) -> Result<Vec<String>, String>;
22    fn view_schema(&self, table: &str) -> Result<Option<RowSchema>, SQLError>;
23}
24/// Capture names and types from an active scope without exposing execution state.
25pub trait ReturningScope {
26    fn binding_snapshot(&self) -> Result<BindingSnapshot, SQLError>;
27}
28#[derive(Clone, Copy)]
29pub struct ReturningAnalysisContext<'a> {
30    pub catalog: &'a dyn ReturningCatalog,
31    pub routines: &'a dyn RoutineResolution,
32    pub aggregates: &'a dyn AggregateClassifier,
33    pub scope: &'a dyn ReturningScope,
34}
35fn dml_storage_error(action: &str, error: impl std::fmt::Display) -> SQLError {
36    SQLError::Internal(format!("{action} failed in storage backend: {error}"))
37}
38
39pub fn returning_target_schema(
40    catalog: &dyn ReturningCatalog,
41    table: &str,
42) -> Result<RowSchema, SQLError> {
43    let definitions = catalog
44        .try_describe_table_row_type(table)
45        .map_err(|error| dml_storage_error("RETURNING schema lookup", error))?;
46    let Some(definitions) = definitions else {
47        return catalog
48            .view_schema(table)?
49            .ok_or_else(|| SQLError::UnknownTable(table.to_string()));
50    };
51    if definitions.is_empty() {
52        let columns = catalog
53            .try_table_columns(table)
54            .map_err(|error| dml_storage_error("RETURNING schema lookup", error))?;
55        let width = columns.len();
56        return Ok(RowSchema::with_types(columns, vec![None; width]));
57    }
58    let columns = definitions
59        .iter()
60        .map(|definition| definition.name.clone())
61        .collect();
62    let types = definitions
63        .into_iter()
64        .map(|definition| Some(definition.ty))
65        .collect();
66    Ok(RowSchema::with_types(columns, types))
67}
68
69pub fn expanded_returning_projections(
70    catalog: &dyn ReturningCatalog,
71    table: &str,
72    target_qualifier: &str,
73    aliases: &ReturningAliases,
74    returning: &[ProjectionPlan],
75) -> Result<Vec<ProjectionPlan>, SQLError> {
76    let columns = returning_target_schema(catalog, table)?.columns().to_vec();
77    let mut projections = Vec::with_capacity(returning.len().max(columns.len()));
78    for projection in returning {
79        match &projection.expr {
80            crate::ScalarExpr::Star => {
81                projections.extend(columns.iter().map(|column| ProjectionPlan {
82                    expr: crate::ScalarExpr::Column(column.clone()),
83                    alias: Some(column.clone()),
84                }));
85            }
86            crate::ScalarExpr::QualifiedStar(qualifier)
87                if qualifier == target_qualifier
88                    || qualifier == &aliases.old
89                    || qualifier == &aliases.new =>
90            {
91                projections.extend(columns.iter().map(|column| ProjectionPlan {
92                    expr: crate::ScalarExpr::QualifiedColumn {
93                        qualifier: qualifier.clone(),
94                        column: column.clone(),
95                    },
96                    alias: Some(column.clone()),
97                }));
98            }
99            _ => projections.push(projection.clone()),
100        }
101    }
102    if !returning.is_empty() && projections.is_empty() {
103        return Err(SQLError::Routine {
104            sqlstate: "42601".into(),
105            message: "RETURNING must have at least one column".into(),
106        });
107    }
108    Ok(projections)
109}
110
111pub fn dml_statement_returning_schema(
112    context: ReturningAnalysisContext<'_>,
113    statement: Statement,
114) -> Result<Option<RowSchema>, SQLError> {
115    let plan = crate::plan::UnifiedPlan::lower_with(statement, context.aggregates);
116    let crate::plan::UnifiedPlan::Command(command) = plan else {
117        return Ok(None);
118    };
119    dml_command_returning_schema(context, &command, &[])
120}
121
122pub fn dml_command_returning_schema(
123    context: ReturningAnalysisContext<'_>,
124    command: &crate::plan::CommandPlan,
125    params: &[SQLParam],
126) -> Result<Option<RowSchema>, SQLError> {
127    match command {
128        crate::plan::CommandPlan::Insert(plan) => analyze_dml_returning_plan(
129            context,
130            &plan.table,
131            &plan.target_qualifier,
132            &plan.returning_aliases,
133            &plan.returning,
134            &plan.ctes,
135            None,
136            &plan.subqueries,
137            params,
138        ),
139        crate::plan::CommandPlan::Update(plan) => analyze_dml_returning_plan(
140            context,
141            &plan.table,
142            &plan.target_qualifier,
143            &plan.returning_aliases,
144            &plan.returning,
145            &plan.ctes,
146            plan.source.as_deref(),
147            &plan.subqueries,
148            params,
149        ),
150        crate::plan::CommandPlan::Delete(plan) => analyze_dml_returning_plan(
151            context,
152            &plan.table,
153            &plan.target_qualifier,
154            &plan.returning_aliases,
155            &plan.returning,
156            &plan.ctes,
157            plan.source.as_deref(),
158            &plan.subqueries,
159            params,
160        ),
161        _ => Ok(None),
162    }
163}
164
165pub fn validate_insert_returning(
166    context: ReturningAnalysisContext<'_>,
167    plan: &crate::plan::InsertPlan,
168    params: &[SQLParam],
169    inherited: Option<&dyn ReturningScope>,
170) -> Result<(), SQLError> {
171    if plan.returning.is_empty() {
172        return Ok(());
173    }
174    let target = returning_target_schema(context.catalog, &plan.table)?;
175    if plan.source.is_none() {
176        let width = if plan.columns.is_empty() {
177            target.len()
178        } else {
179            plan.columns.len()
180        };
181        for row in &plan.rows {
182            if row.len() > width || (!plan.columns.is_empty() && row.len() < width) {
183                return Err(SQLError::Routine {
184                    sqlstate: "42601".into(),
185                    message: if row.len() > width {
186                        "INSERT has more expressions than target columns"
187                    } else {
188                        "INSERT has more target columns than expressions"
189                    }
190                    .into(),
191                });
192            }
193        }
194    }
195    let mut scope = inherited.unwrap_or(context.scope).binding_snapshot()?;
196    for cte in &plan.ctes {
197        scope.insert_deferred(cte.clone());
198    }
199    scope.scalar_subqueries.clone_from(&plan.subqueries);
200    let expressions = returning_expression_schema(
201        &target,
202        &plan.target_qualifier,
203        &plan.returning_aliases,
204        None,
205    );
206    let projections = expanded_returning_projections(
207        context.catalog,
208        &plan.table,
209        &plan.target_qualifier,
210        &plan.returning_aliases,
211        &plan.returning,
212    )?;
213    crate::binding::analyze_projection_output_schema(
214        context.routines,
215        &projections,
216        &expressions,
217        &target,
218        &plan.subqueries,
219        params,
220        &scope.context(),
221    )?;
222    Ok(())
223}
224
225#[expect(
226    clippy::too_many_arguments,
227    reason = "keeps RETURNING row types and nested scopes explicit"
228)]
229pub fn analyze_dml_returning_plan(
230    context: ReturningAnalysisContext<'_>,
231    table: &str,
232    target_qualifier: &str,
233    aliases: &ReturningAliases,
234    returning: &[ProjectionPlan],
235    cte_plans: &[crate::plan::CtePlan],
236    source: Option<&crate::plan::SourcePlan>,
237    subqueries: &[crate::plan::QueryPlan],
238    params: &[SQLParam],
239) -> Result<Option<RowSchema>, SQLError> {
240    if returning.is_empty() {
241        return Ok(None);
242    }
243    let mut ctes = context.scope.binding_snapshot()?;
244    for plan in cte_plans {
245        ctes.insert_deferred(plan.clone());
246    }
247    ctes.scalar_subqueries = subqueries.to_vec();
248    let supplemental = source
249        .map(|source| {
250            crate::binding::analyze_source_plan_schema(
251                context.routines,
252                source,
253                params,
254                &ctes.context(),
255                None,
256            )
257        })
258        .transpose()?;
259    let star_schema = returning_target_schema(context.catalog, table)?;
260    let expression_schema = returning_expression_schema(
261        &star_schema,
262        target_qualifier,
263        aliases,
264        supplemental.as_ref(),
265    );
266    let projections = expanded_returning_projections(
267        context.catalog,
268        table,
269        target_qualifier,
270        aliases,
271        returning,
272    )?;
273    crate::binding::analyze_projection_output_schema(
274        context.routines,
275        &projections,
276        &expression_schema,
277        &star_schema,
278        subqueries,
279        params,
280        &ctes.context(),
281    )
282    .map(Some)
283}
284
285pub fn document_supplied_id(
286    document: &Document,
287    id_column: &str,
288    auto_increment: bool,
289) -> Result<Option<DocId>, SQLError> {
290    match document.get(id_column) {
291        Some(Value::Int(value)) if *value >= 0 => Ok(Some(*value as DocId)),
292        Some(Value::Null) | None => Ok(None),
293        Some(other) if auto_increment => Err(SQLError::TypeMismatch(format!(
294            "auto-increment id must be an integer, got {other:?}"
295        ))),
296        Some(_) => Ok(None),
297    }
298}
299
300pub fn validate_returning_alias_relations(
301    target_qualifier: &str,
302    aliases: &ReturningAliases,
303    supplemental: Option<&RowSchema>,
304) -> Result<(), SQLError> {
305    let mut relation_names = BTreeSet::from([target_qualifier]);
306    for (alias, explicit) in [
307        (aliases.old.as_str(), aliases.old_explicit),
308        (aliases.new.as_str(), aliases.new_explicit),
309    ] {
310        if !explicit {
311            continue;
312        }
313        if relation_names.contains(alias)
314            || supplemental.is_some_and(|schema| schema.has_qualifier(alias))
315        {
316            return Err(SQLError::Routine {
317                sqlstate: "42712".into(),
318                message: format!("table name \"{alias}\" specified more than once"),
319            });
320        }
321        relation_names.insert(alias);
322    }
323    Ok(())
324}