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::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}
24pub 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}