uqa_sql/binding/
commands.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Static RETURNING schemas for command CTEs.
8
9use super::{
10    projection_columns, projection_star_columns, ColumnType, RowSchema, SQLError, SQLParam,
11    ScalarExpr, SchemaScope, SourcePlan,
12};
13use crate::plan::{CommandPlan, CtePlanBody};
14use crate::routines::RoutineResolution;
15
16pub fn analyze_prepared_command_schema(
17    routines: &dyn RoutineResolution,
18    command: &CommandPlan,
19    params: &[SQLParam],
20    ctes: &super::BindingContext,
21) -> Result<Option<RowSchema>, SQLError> {
22    if command.mutation_target().is_none() {
23        return Ok(None);
24    }
25    let mut scope = SchemaScope::for_analysis(ctes)?;
26    scope.set_command_lookup_mode(command);
27    let previous = scope.bind_cte_schemas(routines, command.ctes(), params, None)?;
28    let result = (|| {
29        let (target, expression) = scope.command_expression_schema(routines, command, params)?;
30        let excluded = RowSchema::with_qualified_types(
31            "excluded",
32            target.columns().to_vec(),
33            target.column_types().to_vec(),
34        );
35        let conflict_input = RowSchema::join(&expression, &excluded, std::iter::empty::<String>());
36        let conflict_expressions = match command {
37            CommandPlan::Insert(insert) => insert
38                .on_conflict
39                .as_ref()
40                .map(|conflict| match &conflict.action {
41                    crate::plan::ConflictActionPlan::Update {
42                        assignments,
43                        predicate,
44                    } => assignments
45                        .iter()
46                        .map(|assignment| &assignment.value)
47                        .chain(predicate.as_deref())
48                        .collect::<Vec<_>>(),
49                    crate::plan::ConflictActionPlan::Nothing => Vec::new(),
50                })
51                .unwrap_or_default(),
52            _ => Vec::new(),
53        };
54        if let CommandPlan::Insert(insert) = command {
55            if let Some(source) = &insert.source {
56                scope.bind_query(routines, source, params, None)?;
57            }
58        }
59        for scalar in command.expressions() {
60            if matches!(command, CommandPlan::Merge(_)) {
61                continue;
62            }
63            scope.bind_expression_type(
64                routines,
65                scalar,
66                if conflict_expressions
67                    .iter()
68                    .any(|candidate| std::ptr::eq(*candidate, scalar))
69                {
70                    &conflict_input
71                } else {
72                    &expression
73                },
74                command.scalar_subqueries(),
75                params,
76                None,
77            )?;
78        }
79        if let CommandPlan::Merge(merge) = command {
80            scope.bind_merge_expressions(routines, merge, params, &target, &expression)?;
81        }
82        let result = scope.bind_command_returning(routines, command, params)?;
83        Ok(command
84            .returning()
85            .filter(|returning| !returning.is_empty())
86            .map(|_| result))
87    })();
88    scope.restore_cte_schemas(previous);
89    result
90}
91
92impl SchemaScope {
93    pub(super) fn set_command_lookup_mode(&mut self, command: &CommandPlan) {
94        let bound = match command {
95            CommandPlan::Insert(plan) => plan.relations_bound,
96            CommandPlan::Update(plan) => plan.relations_bound,
97            CommandPlan::Delete(plan) => plan.relations_bound,
98            _ => false,
99        };
100        self.resolution.set_lookup_mode(if bound {
101            crate::catalog::resolution::RelationLookupMode::Bound
102        } else {
103            crate::catalog::resolution::RelationLookupMode::Dynamic
104        });
105    }
106
107    pub(super) fn bind_cte_body(
108        &mut self,
109        routines: &dyn RoutineResolution,
110        body: &CtePlanBody,
111        params: &[SQLParam],
112        outer: Option<&RowSchema>,
113    ) -> Result<RowSchema, SQLError> {
114        match body {
115            CtePlanBody::Query(query) => self.bind_query(routines, query, params, outer),
116            CtePlanBody::Command(command) => self.bind_command_returning(routines, command, params),
117        }
118    }
119
120    fn bind_command_target(
121        &mut self,
122        routines: &dyn RoutineResolution,
123        command: &CommandPlan,
124        params: &[SQLParam],
125    ) -> Result<RowSchema, SQLError> {
126        let table = command
127            .mutation_target()
128            .ok_or_else(|| SQLError::Internal("non-DML command in a WITH definition".into()))?;
129        let source = SourcePlan::Table {
130            bound_columns: None,
131            name: table.to_string(),
132            qualifier: command.target_qualifier().unwrap_or(table).to_string(),
133            alias: None,
134            column_aliases: Vec::new(),
135            include_descendants: true,
136        };
137        // A mutation target always denotes a catalog relation, even when a CTE shadows its name.
138        let previous = std::mem::take(&mut self.ctes);
139        let deferred = std::mem::take(&mut self.deferred_ctes);
140        let bound = match command {
141            CommandPlan::Insert(plan) => plan.target_relation_bound,
142            CommandPlan::Update(plan) => plan.target_relation_bound,
143            CommandPlan::Delete(plan) => plan.target_relation_bound,
144            _ => false,
145        };
146        let lookup = bound.then(|| {
147            self.resolution
148                .set_lookup_mode(crate::catalog::resolution::RelationLookupMode::Bound)
149        });
150        let result = self.bind_source(routines, &source, &[], params, None);
151        if let Some(lookup) = lookup {
152            self.resolution.set_lookup_mode(lookup);
153        }
154        self.ctes = previous;
155        self.deferred_ctes = deferred;
156        let target = result?;
157        let columns = target
158            .columns()
159            .iter()
160            .enumerate()
161            .map(|(position, column)| target.public_name(position).unwrap_or(column).to_string())
162            .collect();
163        let schema = RowSchema::with_types(columns, target.column_types().to_vec());
164        Ok(if target.columns_are_open(None) {
165            RowSchema::with_open_columns(&schema, None)
166        } else {
167            schema
168        })
169    }
170
171    pub(super) fn command_expression_schema(
172        &mut self,
173        routines: &dyn RoutineResolution,
174        command: &CommandPlan,
175        params: &[SQLParam],
176    ) -> Result<(RowSchema, RowSchema), SQLError> {
177        let target = self.bind_command_target(routines, command, params)?;
178        let source = command
179            .source_input()
180            .map(|source| {
181                self.bind_source(routines, source, command.scalar_subqueries(), params, None)
182            })
183            .transpose()?;
184        let aliases = command
185            .returning_aliases()
186            .ok_or_else(|| SQLError::Internal("command CTE has no RETURNING namespace".into()))?;
187        let mut expression = crate::semantics::returning_expression_schema(
188            &target,
189            command.target_qualifier().unwrap_or_default(),
190            aliases,
191            source.as_ref(),
192        );
193        if target.columns_are_open(None) {
194            expression = RowSchema::with_open_columns(&expression, command.target_qualifier());
195        }
196        Ok((target, expression))
197    }
198
199    pub(super) fn bind_command_returning(
200        &mut self,
201        routines: &dyn RoutineResolution,
202        command: &CommandPlan,
203        params: &[SQLParam],
204    ) -> Result<RowSchema, SQLError> {
205        let previous = self.bind_cte_schemas(routines, command.ctes(), params, None)?;
206        let result = (|| {
207            let (target, expression) = self.command_expression_schema(routines, command, params)?;
208            if let CommandPlan::Insert(insert) = command {
209                if let Some(source) = &insert.source {
210                    self.bind_query(routines, source, params, None)?;
211                }
212            }
213            let returning = command.returning().unwrap_or_default();
214            let labels = projection_columns(returning);
215            let mut columns = Vec::new();
216            let mut types: Vec<Option<ColumnType>> = Vec::new();
217            for (position, projection) in returning.iter().enumerate() {
218                let expansion = if matches!(projection.expr, ScalarExpr::QualifiedStar(_)) {
219                    &expression
220                } else {
221                    &target
222                };
223                if let Some(star) = projection_star_columns(&projection.expr, expansion)? {
224                    for (column, ty) in star {
225                        columns.push(column);
226                        types.push(ty);
227                    }
228                } else {
229                    columns.push(labels[position].clone());
230                    types.push(self.bind_expression_type(
231                        routines,
232                        &projection.expr,
233                        &expression,
234                        command.scalar_subqueries(),
235                        params,
236                        Some(&expression),
237                    )?);
238                }
239            }
240            if let Some(error) =
241                crate::semantics::virtual_relation_mutation_error(&self.resolution, command)
242            {
243                return Err(error);
244            }
245            Ok(RowSchema::with_types(columns, types))
246        })();
247        self.restore_cte_schemas(previous);
248        result
249    }
250}
uqa_sql/binding/commands.rs

uqa_sql/binding/
commands.rs
