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::{borrow::Cow, collections::BTreeSet};
17use uqa_core::{DocId, Value};
18
19#[cfg(test)]
20mod tests;
21
22pub trait ReturningCatalog {
23    fn try_describe_table_row_type(&self, table: &str) -> Result<Option<Vec<ColumnDef>>, String>;
24    fn try_table_columns(&self, table: &str) -> Result<Vec<String>, String>;
25    fn view_schema(&self, table: &str) -> Result<Option<RowSchema>, SQLError>;
26}
27/// Capture names and types from an active scope without exposing execution state.
28pub trait ReturningScope {
29    fn binding_snapshot(&self) -> Result<BindingSnapshot, SQLError>;
30}
31#[derive(Clone, Copy)]
32pub struct ReturningAnalysisContext<'a> {
33    pub catalog: &'a dyn ReturningCatalog,
34    pub routines: &'a dyn RoutineResolution,
35    pub aggregates: &'a dyn AggregateClassifier,
36    pub scope: &'a dyn ReturningScope,
37}
38fn dml_storage_error(action: &str, error: impl std::fmt::Display) -> SQLError {
39    SQLError::Internal(format!("{action} failed in storage backend: {error}"))
40}
41
42pub fn returning_target_schema(
43    catalog: &dyn ReturningCatalog,
44    table: &str,
45) -> Result<RowSchema, SQLError> {
46    let definitions = catalog
47        .try_describe_table_row_type(table)
48        .map_err(|error| dml_storage_error("RETURNING schema lookup", error))?;
49    let Some(definitions) = definitions else {
50        return catalog
51            .view_schema(table)?
52            .ok_or_else(|| SQLError::UnknownTable(table.to_string()));
53    };
54    if definitions.is_empty() {
55        let columns = catalog
56            .try_table_columns(table)
57            .map_err(|error| dml_storage_error("RETURNING schema lookup", error))?;
58        let width = columns.len();
59        return Ok(RowSchema::with_types(columns, vec![None; width]));
60    }
61    Ok(RowSchema::with_column_metadata(
62        None,
63        definitions
64            .into_iter()
65            .map(|column| (column.name, Some(column.ty), column.collation_oid)),
66    ))
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 target = returning_target_schema(catalog, table)?;
77    expand_returning_target_stars(target.columns(), target_qualifier, aliases, returning, true)
78        .map(|projections| projections.into_iter().map(Cow::into_owned).collect())
79}
80
81pub(crate) fn command_returning_projections<'a>(
82    command: &'a crate::plan::CommandPlan,
83    target: &RowSchema,
84) -> Result<Vec<Cow<'a, ProjectionPlan>>, SQLError> {
85    let aliases = command
86        .returning_aliases()
87        .ok_or_else(|| SQLError::Internal("command has no RETURNING namespace".into()))?;
88    expand_returning_target_stars(
89        target.columns(),
90        command.target_qualifier().unwrap_or_default(),
91        aliases,
92        command.returning().unwrap_or_default(),
93        !matches!(command, crate::plan::CommandPlan::Merge(_)),
94    )
95}
96
97pub(crate) fn command_returning_star_schema(
98    command: &crate::plan::CommandPlan,
99    target: &RowSchema,
100    source: Option<&RowSchema>,
101) -> Result<RowSchema, SQLError> {
102    if matches!(command, crate::plan::CommandPlan::Merge(_)) {
103        let source = source
104            .ok_or_else(|| SQLError::Internal("MERGE has no declared source scope".into()))?;
105        Ok(RowSchema::join(source, target, std::iter::empty()))
106    } else {
107        Ok(target.clone())
108    }
109}
110
111pub(crate) fn validate_returning_width(
112    returning: &[ProjectionPlan],
113    width: usize,
114) -> Result<(), SQLError> {
115    if !returning.is_empty() && width == 0 {
116        return Err(SQLError::Routine {
117            sqlstate: "42601".into(),
118            message: "RETURNING must have at least one column".into(),
119        });
120    }
121    Ok(())
122}
123
124/// All RETURNING consumers expand target and image wildcards from the declared user columns. System columns remain individually addressable but are never wildcard outputs.
125fn expand_returning_target_stars<'a>(
126    columns: &[String],
127    target_qualifier: &str,
128    aliases: &ReturningAliases,
129    returning: &'a [ProjectionPlan],
130    expand_bare_star: bool,
131) -> Result<Vec<Cow<'a, ProjectionPlan>>, SQLError> {
132    let mut projections = Vec::with_capacity(returning.len().max(columns.len()));
133    for projection in returning {
134        match &projection.expr {
135            crate::ScalarExpr::Star if expand_bare_star => {
136                projections.extend(columns.iter().map(|column| {
137                    Cow::Owned(ProjectionPlan {
138                        expr: crate::ScalarExpr::qualified_column(target_qualifier, column),
139                        alias: Some(column.clone()),
140                    })
141                }));
142            }
143            crate::ScalarExpr::QualifiedStar(qualifier)
144                if qualifier == target_qualifier
145                    || qualifier == &aliases.old
146                    || qualifier == &aliases.new =>
147            {
148                projections.extend(columns.iter().map(|column| {
149                    Cow::Owned(ProjectionPlan {
150                        expr: crate::ScalarExpr::qualified_column(qualifier, column),
151                        alias: Some(column.clone()),
152                    })
153                }));
154            }
155            _ => projections.push(Cow::Borrowed(projection)),
156        }
157    }
158    validate_returning_width(returning, projections.len())?;
159    Ok(projections)
160}
161
162pub fn dml_statement_returning_schema(
163    context: ReturningAnalysisContext<'_>,
164    statement: Statement,
165) -> Result<Option<RowSchema>, SQLError> {
166    let plan = crate::plan::UnifiedPlan::lower_with(statement, context.aggregates);
167    let crate::plan::UnifiedPlan::Command(command) = plan else {
168        return Ok(None);
169    };
170    dml_command_returning_schema(context, &command, &[])
171}
172
173pub fn dml_command_returning_schema(
174    context: ReturningAnalysisContext<'_>,
175    command: &crate::plan::CommandPlan,
176    params: &[SQLParam],
177) -> Result<Option<RowSchema>, SQLError> {
178    match command {
179        crate::plan::CommandPlan::Insert(plan) => analyze_dml_returning_plan(
180            context,
181            &plan.table,
182            &plan.target_qualifier,
183            &plan.returning_aliases,
184            &plan.returning,
185            &plan.ctes,
186            None,
187            &plan.subqueries,
188            params,
189        ),
190        crate::plan::CommandPlan::Update(plan) => analyze_dml_returning_plan(
191            context,
192            &plan.table,
193            &plan.target_qualifier,
194            &plan.returning_aliases,
195            &plan.returning,
196            &plan.ctes,
197            plan.source.as_deref(),
198            &plan.subqueries,
199            params,
200        ),
201        crate::plan::CommandPlan::Delete(plan) => analyze_dml_returning_plan(
202            context,
203            &plan.table,
204            &plan.target_qualifier,
205            &plan.returning_aliases,
206            &plan.returning,
207            &plan.ctes,
208            plan.source.as_deref(),
209            &plan.subqueries,
210            params,
211        ),
212        _ => Ok(None),
213    }
214}
215
216pub fn validate_insert_returning(
217    context: ReturningAnalysisContext<'_>,
218    plan: &crate::plan::InsertPlan,
219    params: &[SQLParam],
220    inherited: Option<&dyn ReturningScope>,
221) -> Result<(), SQLError> {
222    if plan.returning.is_empty() {
223        return Ok(());
224    }
225    let target = returning_target_schema(context.catalog, &plan.table)?;
226    if plan.source.is_none() {
227        let width = if plan.columns.is_empty() {
228            target.len()
229        } else {
230            plan.columns.len()
231        };
232        for row in &plan.rows {
233            if row.len() > width || (!plan.columns.is_empty() && row.len() < width) {
234                return Err(SQLError::Routine {
235                    sqlstate: "42601".into(),
236                    message: if row.len() > width {
237                        "INSERT has more expressions than target columns"
238                    } else {
239                        "INSERT has more target columns than expressions"
240                    }
241                    .into(),
242                });
243            }
244        }
245    }
246    let mut scope = inherited.unwrap_or(context.scope).binding_snapshot()?;
247    for cte in &plan.ctes {
248        scope.insert_deferred(cte.clone());
249    }
250    scope.scalar_subqueries.clone_from(&plan.subqueries);
251    let expressions = returning_expression_schema(
252        &target,
253        &plan.target_qualifier,
254        &plan.returning_aliases,
255        None,
256    );
257    let projections = expanded_returning_projections(
258        context.catalog,
259        &plan.table,
260        &plan.target_qualifier,
261        &plan.returning_aliases,
262        &plan.returning,
263    )?;
264    crate::binding::analyze_projection_output_schema(
265        context.routines,
266        &projections,
267        &expressions,
268        &target,
269        &plan.subqueries,
270        params,
271        &scope.context(),
272    )?;
273    Ok(())
274}
275
276#[expect(
277    clippy::too_many_arguments,
278    reason = "keeps RETURNING row types and nested scopes explicit"
279)]
280pub fn analyze_dml_returning_plan(
281    context: ReturningAnalysisContext<'_>,
282    table: &str,
283    target_qualifier: &str,
284    aliases: &ReturningAliases,
285    returning: &[ProjectionPlan],
286    cte_plans: &[crate::plan::CtePlan],
287    source: Option<&crate::plan::SourcePlan>,
288    subqueries: &[crate::plan::QueryPlan],
289    params: &[SQLParam],
290) -> Result<Option<RowSchema>, SQLError> {
291    if returning.is_empty() {
292        return Ok(None);
293    }
294    let mut ctes = context.scope.binding_snapshot()?;
295    for plan in cte_plans {
296        ctes.insert_deferred(plan.clone());
297    }
298    ctes.scalar_subqueries = subqueries.to_vec();
299    let supplemental = source
300        .map(|source| {
301            crate::binding::analyze_source_plan_schema(
302                context.routines,
303                source,
304                params,
305                &ctes.context(),
306                None,
307            )
308        })
309        .transpose()?;
310    let star_schema = returning_target_schema(context.catalog, table)?;
311    let expression_schema = returning_expression_schema(
312        &star_schema,
313        target_qualifier,
314        aliases,
315        supplemental.as_ref(),
316    );
317    let projections = expanded_returning_projections(
318        context.catalog,
319        table,
320        target_qualifier,
321        aliases,
322        returning,
323    )?;
324    crate::binding::analyze_projection_output_schema(
325        context.routines,
326        &projections,
327        &expression_schema,
328        &star_schema,
329        subqueries,
330        params,
331        &ctes.context(),
332    )
333    .map(Some)
334}
335
336/// The identity the `id` field of a document without declared columns names, as the document API addresses it: any nonnegative integer.
337pub fn document_supplied_id(document: &Document, id_column: &str) -> Option<DocId> {
338    match document.get(id_column) {
339        Some(Value::Int(value)) => DocId::try_from(*value).ok(),
340        _ => None,
341    }
342}
343
344pub fn validate_returning_alias_relations(
345    target_qualifier: &str,
346    aliases: &ReturningAliases,
347    supplemental: Option<&RowSchema>,
348) -> Result<(), SQLError> {
349    let mut relation_names = BTreeSet::from([target_qualifier]);
350    for (alias, explicit) in [
351        (aliases.old.as_str(), aliases.old_explicit),
352        (aliases.new.as_str(), aliases.new_explicit),
353    ] {
354        if !explicit {
355            continue;
356        }
357        if relation_names.contains(alias)
358            || supplemental.is_some_and(|schema| schema.has_qualifier(alias))
359        {
360            return Err(SQLError::Routine {
361                sqlstate: "42712".into(),
362                message: format!("table name \"{alias}\" specified more than once"),
363            });
364        }
365        relation_names.insert(alias);
366    }
367    Ok(())
368}