1use 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}
27pub 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
124fn 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
336pub 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}