Skip to main content

uqa_sql/semantics/
merge.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! MERGE clause scopes, target privileges, and positional RETURNING schemas.
8use super::returning::expanded_returning_projections;
9use crate::{
10    plan::{MergePlan, MergeWhenPlan, ProjectionPlan},
11    SQLError, SQLParam,
12};
13use std::collections::BTreeSet;
14
15#[expect(
16    clippy::too_many_lines,
17    reason = "validates clause-specific schemas and FULL JOIN requirements"
18)]
19pub fn validate_merge_action_scopes(
20    routines: &dyn crate::routines::RoutineResolution,
21    stmt: &MergePlan,
22    target_schema: &crate::RowSchema,
23    source_schema: &crate::RowSchema,
24    params: &[SQLParam],
25    bindings: &crate::binding::context::BindingContext<'_>,
26) -> Result<(), SQLError> {
27    let matched_schema = crate::RowSchema::join(target_schema, source_schema, std::iter::empty());
28    let expression_type = |expression: &crate::ScalarExpr, schema: &crate::RowSchema| {
29        crate::binding::analyze_projection_output_schema(
30            routines,
31            &[ProjectionPlan {
32                expr: expression.clone(),
33                alias: None,
34            }],
35            schema,
36            schema,
37            &stmt.subqueries,
38            params,
39            bindings,
40        )
41        .map(|output| output.column_type(0).cloned())
42    };
43    let validate_boolean = |expression: &crate::ScalarExpr,
44                            schema: &crate::RowSchema,
45                            label: &str|
46     -> Result<(), SQLError> {
47        if expression_type(expression, schema)?
48            .is_some_and(|ty| ty != crate::ast::ColumnType::Boolean)
49        {
50            return Err(SQLError::TypeMismatch(format!(
51                "argument of {label} must be type boolean"
52            )));
53        }
54        Ok(())
55    };
56    validate_boolean(&stmt.join_condition, &matched_schema, "MERGE ON")?;
57    let has_source_missing = stmt.when_clauses.iter().any(|clause| {
58        matches!(
59            clause,
60            MergeWhenPlan::UpdateNotMatchedBySource { .. }
61                | MergeWhenPlan::DeleteNotMatchedBySource { .. }
62                | MergeWhenPlan::NothingNotMatchedBySource { .. }
63        )
64    });
65    let has_target_missing = stmt.when_clauses.iter().any(|clause| {
66        matches!(
67            clause,
68            MergeWhenPlan::InsertNotMatched { .. } | MergeWhenPlan::NothingNotMatched { .. }
69        )
70    });
71    if has_source_missing
72        && has_target_missing
73        && !super::join_predicates::join_conjuncts(&stmt.join_condition)
74            .into_iter()
75            .any(|conjunct| {
76                matches!(
77                    conjunct,
78                    crate::ScalarExpr::Binary {
79                        op: crate::ast::BinaryOp::Equal,
80                        lhs,
81                        rhs,
82                    } if super::join_predicates::decide_join_sides(
83                        target_schema,
84                        source_schema,
85                        lhs,
86                        rhs,
87                    )
88                    .is_some()
89                )
90            })
91    {
92        return Err(SQLError::Routine {
93            sqlstate: "0A000".into(),
94            message:
95                "FULL JOIN is only supported with merge-joinable or hash-joinable join conditions"
96                    .into(),
97        });
98    }
99    for clause in &stmt.when_clauses {
100        let (condition, expressions, schema): (
101            Option<&crate::ScalarExpr>,
102            Vec<&crate::ScalarExpr>,
103            &crate::RowSchema,
104        ) = match clause {
105            MergeWhenPlan::UpdateMatched {
106                condition,
107                assignments,
108            } => (
109                condition.as_ref(),
110                assignments
111                    .iter()
112                    .flat_map(crate::plan::AssignmentPlan::expressions)
113                    .collect(),
114                &matched_schema,
115            ),
116            MergeWhenPlan::DeleteMatched { condition }
117            | MergeWhenPlan::NothingMatched { condition } => {
118                (condition.as_ref(), Vec::new(), &matched_schema)
119            }
120            MergeWhenPlan::UpdateNotMatchedBySource {
121                condition,
122                assignments,
123            } => (
124                condition.as_ref(),
125                assignments
126                    .iter()
127                    .flat_map(crate::plan::AssignmentPlan::expressions)
128                    .collect(),
129                target_schema,
130            ),
131            MergeWhenPlan::DeleteNotMatchedBySource { condition }
132            | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
133                (condition.as_ref(), Vec::new(), target_schema)
134            }
135            MergeWhenPlan::InsertNotMatched {
136                condition,
137                columns,
138                values,
139            } => (
140                condition.as_ref(),
141                columns
142                    .iter()
143                    .flat_map(crate::ast::AssignmentTarget::expressions)
144                    .chain(values)
145                    .collect(),
146                source_schema,
147            ),
148            MergeWhenPlan::NothingNotMatched { condition } => {
149                (condition.as_ref(), Vec::new(), source_schema)
150            }
151        };
152        if let Some(condition) = condition {
153            validate_boolean(condition, schema, "WHEN")?;
154        }
155        for expression in expressions {
156            expression_type(expression, schema)?;
157        }
158    }
159    Ok(())
160}
161
162pub fn expanded_merge_returning_projections(
163    catalog: &dyn super::returning::ReturningCatalog,
164    target_table: &str,
165    target_qualifier: &str,
166    aliases: &crate::ast::ReturningAliases,
167    source_schema: &crate::RowSchema,
168    source_relation: crate::ast::InternalRelationId,
169    returning: &[ProjectionPlan],
170) -> Result<Vec<ProjectionPlan>, SQLError> {
171    let target_star = ProjectionPlan {
172        expr: crate::ScalarExpr::QualifiedStar(target_qualifier.into()),
173        alias: None,
174    };
175    let target_projections = expanded_returning_projections(
176        catalog,
177        target_table,
178        target_qualifier,
179        aliases,
180        std::slice::from_ref(&target_star),
181    )?;
182    let mut projections = Vec::new();
183    for projection in returning {
184        match &projection.expr {
185            crate::ScalarExpr::Star => {
186                projections.extend(
187                    source_schema
188                        .columns()
189                        .iter()
190                        .enumerate()
191                        .filter(|(position, _)| {
192                            super::projection::visible_projection_source_position(
193                                source_schema,
194                                *position,
195                            )
196                        })
197                        .map(|(position, column)| ProjectionPlan {
198                            expr: crate::ScalarExpr::InternalColumn(
199                                source_relation.column(position),
200                            ),
201                            alias: Some(
202                                source_schema
203                                    .public_name(position)
204                                    .unwrap_or(column)
205                                    .to_string(),
206                            ),
207                        }),
208                );
209                projections.extend(target_projections.iter().cloned());
210            }
211            crate::ScalarExpr::QualifiedStar(qualifier)
212                if qualifier == target_qualifier
213                    || qualifier == &aliases.old
214                    || qualifier == &aliases.new =>
215            {
216                projections.extend(expanded_returning_projections(
217                    catalog,
218                    target_table,
219                    target_qualifier,
220                    aliases,
221                    std::slice::from_ref(projection),
222                )?);
223            }
224            _ => projections.push(projection.clone()),
225        }
226    }
227    Ok(projections)
228}
229
230pub fn merge_returning_source_schema(
231    source_schema: &crate::RowSchema,
232    source_relation: crate::ast::InternalRelationId,
233) -> crate::RowSchema {
234    let aliases = source_schema
235        .columns()
236        .iter()
237        .enumerate()
238        .map(|(position, _)| {
239            (
240                source_relation.column(position),
241                source_schema
242                    .physical_slot(position)
243                    .expect("source column has a physical slot"),
244                source_schema.column_type(position).cloned(),
245            )
246        })
247        .collect::<Vec<_>>();
248    crate::RowSchema::with_physical_internal_aliases(source_schema, &aliases)
249}
250
251pub fn ensure_merge_mutation_privileges(
252    catalog: &dyn super::mutation_privileges::MutationPrivilegeCatalog,
253    stmt: &MergePlan,
254) -> Result<(), SQLError> {
255    let mut column_privileges = BTreeSet::new();
256    let mut requires_delete = false;
257    let mut requires_any_insert = false;
258    let table_columns = catalog.bound_table_column_names(&stmt.target)?;
259    let privilege_subject = stmt
260        .target_privilege_subject
261        .clone()
262        .unwrap_or_else(|| catalog.current_role());
263    for clause in &stmt.when_clauses {
264        match clause {
265            MergeWhenPlan::InsertNotMatched {
266                columns, values, ..
267            } => {
268                if columns.is_empty() && values.is_empty() {
269                    requires_any_insert = true;
270                } else {
271                    let columns = if columns.is_empty() {
272                        table_columns
273                            .iter()
274                            .take(values.len())
275                            .cloned()
276                            .collect::<Vec<_>>()
277                    } else {
278                        columns.iter().map(|target| target.column.clone()).collect()
279                    };
280                    column_privileges.extend(columns.into_iter().map(|column| {
281                        (
282                            crate::catalog::security::table::TableAclPrivilege::Insert,
283                            column,
284                        )
285                    }));
286                }
287            }
288            MergeWhenPlan::UpdateMatched { assignments, .. }
289            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => {
290                column_privileges.extend(assignments.iter().map(|assignment| {
291                    (
292                        crate::catalog::security::table::TableAclPrivilege::Update,
293                        assignment.target.column.clone(),
294                    )
295                }));
296            }
297            MergeWhenPlan::DeleteMatched { .. }
298            | MergeWhenPlan::DeleteNotMatchedBySource { .. } => requires_delete = true,
299            _ => {}
300        }
301    }
302    if requires_delete {
303        catalog.ensure_table_privilege_for(
304            &stmt.target,
305            &privilege_subject,
306            crate::catalog::security::table::TableAclPrivilege::Delete,
307        )?;
308    }
309    if requires_any_insert {
310        catalog.ensure_any_column_privilege_for(
311            &stmt.target,
312            &privilege_subject,
313            crate::catalog::security::table::TableAclPrivilege::Insert,
314        )?;
315    }
316    for (privilege, column) in column_privileges {
317        catalog.ensure_column_privilege_for(
318            &stmt.target,
319            &column,
320            &privilege_subject,
321            privilege,
322        )?;
323    }
324    Ok(())
325}
326
327pub fn merge_returning_schema(
328    routines: &dyn crate::routines::RoutineResolution,
329    catalog: &dyn super::returning::ReturningCatalog,
330    stmt: &MergePlan,
331    params: &[SQLParam],
332    source_schema: &crate::RowSchema,
333    ctes: &crate::binding::context::BindingContext<'_>,
334) -> Result<Option<crate::RowSchema>, SQLError> {
335    if stmt.returning.is_empty() {
336        return Ok(None);
337    }
338    let source_relation = crate::ast::InternalRelationId::allocate();
339    let projections = expanded_merge_returning_projections(
340        catalog,
341        &stmt.target,
342        &stmt.target_qualifier,
343        &stmt.returning_aliases,
344        source_schema,
345        source_relation,
346        &stmt.returning,
347    )?;
348    let returning_source_schema = merge_returning_source_schema(source_schema, source_relation);
349    let star_schema = super::returning::returning_target_schema(catalog, &stmt.target)?;
350    let expression_schema = super::returning_expression_schema(
351        &star_schema,
352        &stmt.target_qualifier,
353        &stmt.returning_aliases,
354        Some(&returning_source_schema),
355    );
356    crate::binding::analyze_projection_output_schema(
357        routines,
358        &projections,
359        &expression_schema,
360        &star_schema,
361        &stmt.subqueries,
362        params,
363        ctes,
364    )
365    .map(Some)
366}
367
368pub fn merge_command_returning_schema(
369    routines: &dyn crate::routines::RoutineResolution,
370    catalog: &dyn super::returning::ReturningCatalog,
371    rows: &dyn super::mutation_rows::MutationRowCatalog,
372    stmt: &MergePlan,
373    params: &[SQLParam],
374    bindings: &crate::binding::context::BindingContext<'_>,
375) -> Result<Option<crate::RowSchema>, SQLError> {
376    if stmt.returning.is_empty() {
377        return Ok(None);
378    }
379    let source =
380        crate::binding::analyze_source_plan_schema(routines, &stmt.source, params, bindings, None)?;
381    super::returning::validate_returning_alias_relations(
382        &stmt.target_qualifier,
383        &stmt.returning_aliases,
384        Some(&source),
385    )?;
386    let target =
387        super::mutation_rows::null_target_schema(rows, &stmt.target, &stmt.target_qualifier)?;
388    validate_merge_action_scopes(routines, stmt, &target, &source, params, bindings)?;
389    merge_returning_schema(routines, catalog, stmt, params, &source, bindings)
390}
391
392pub fn validate_merge_target_columns(
393    catalog: &dyn crate::assignment::columns::AssignmentColumnCatalog,
394    stmt: &MergePlan,
395) -> Result<(), SQLError> {
396    use crate::assignment::columns::validate_mutation_targets;
397    for clause in &stmt.when_clauses {
398        match clause {
399            MergeWhenPlan::UpdateMatched { assignments, .. }
400            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => {
401                validate_mutation_targets(
402                    catalog,
403                    &stmt.target,
404                    assignments.iter().map(|assignment| &assignment.target),
405                    "MERGE UPDATE",
406                    false,
407                )?;
408            }
409            MergeWhenPlan::InsertNotMatched { columns, .. } => validate_mutation_targets(
410                catalog,
411                &stmt.target,
412                columns.iter(),
413                "MERGE INSERT",
414                true,
415            )?,
416            _ => {}
417        }
418    }
419    Ok(())
420}