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/// The union of target columns of every MERGE UPDATE action, which `PostgreSQL` uses when selecting generated-column dependencies even if a different action is selected for a particular row.
16#[must_use]
17pub fn merge_update_columns(statement: &MergePlan) -> Vec<String> {
18    statement
19        .when_clauses
20        .iter()
21        .filter_map(|clause| match clause {
22            MergeWhenPlan::UpdateMatched { assignments, .. }
23            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => Some(assignments),
24            _ => None,
25        })
26        .flatten()
27        .flat_map(|assignment| assignment.target.column_names())
28        .map(str::to_owned)
29        .collect::<BTreeSet<_>>()
30        .into_iter()
31        .collect()
32}
33
34#[expect(
35    clippy::too_many_lines,
36    reason = "validates clause-specific schemas and FULL JOIN requirements"
37)]
38pub fn validate_merge_action_scopes(
39    routines: &dyn crate::routines::RoutineResolution,
40    stmt: &MergePlan,
41    target_schema: &crate::RowSchema,
42    source_schema: &crate::RowSchema,
43    params: &[SQLParam],
44    bindings: &crate::binding::context::BindingContext<'_>,
45) -> Result<(), SQLError> {
46    let matched_schema = crate::RowSchema::join(target_schema, source_schema, std::iter::empty());
47    let expression_type = |expression: &crate::ScalarExpr, schema: &crate::RowSchema| {
48        crate::binding::analyze_projection_output_schema(
49            routines,
50            &[ProjectionPlan {
51                expr: expression.clone(),
52                alias: None,
53            }],
54            schema,
55            schema,
56            &stmt.subqueries,
57            params,
58            bindings,
59        )
60        .map(|output| output.column_type(0).cloned())
61    };
62    let validate_boolean = |expression: &crate::ScalarExpr,
63                            schema: &crate::RowSchema,
64                            label: &str|
65     -> Result<(), SQLError> {
66        if expression_type(expression, schema)?
67            .is_some_and(|ty| ty != crate::ast::ColumnType::Boolean)
68        {
69            return Err(SQLError::TypeMismatch(format!(
70                "argument of {label} must be type boolean"
71            )));
72        }
73        Ok(())
74    };
75    validate_boolean(&stmt.join_condition, &matched_schema, "MERGE ON")?;
76    let has_source_missing = stmt.when_clauses.iter().any(|clause| {
77        matches!(
78            clause,
79            MergeWhenPlan::UpdateNotMatchedBySource { .. }
80                | MergeWhenPlan::DeleteNotMatchedBySource { .. }
81                | MergeWhenPlan::NothingNotMatchedBySource { .. }
82        )
83    });
84    let has_target_missing = stmt.when_clauses.iter().any(|clause| {
85        matches!(
86            clause,
87            MergeWhenPlan::InsertNotMatched { .. } | MergeWhenPlan::NothingNotMatched { .. }
88        )
89    });
90    if has_source_missing
91        && has_target_missing
92        && !super::join_predicates::join_conjuncts(&stmt.join_condition)
93            .into_iter()
94            .any(|conjunct| {
95                matches!(
96                    conjunct,
97                    crate::ScalarExpr::Binary {
98                        op: crate::ast::BinaryOp::Equal,
99                        lhs,
100                        rhs,
101                    } if super::join_predicates::decide_join_sides(
102                        target_schema,
103                        source_schema,
104                        lhs,
105                        rhs,
106                    )
107                    .is_some()
108                )
109            })
110    {
111        return Err(SQLError::Routine {
112            sqlstate: "0A000".into(),
113            message:
114                "FULL JOIN is only supported with merge-joinable or hash-joinable join conditions"
115                    .into(),
116        });
117    }
118    for clause in &stmt.when_clauses {
119        let (condition, expressions, schema): (
120            Option<&crate::ScalarExpr>,
121            Vec<&crate::ScalarExpr>,
122            &crate::RowSchema,
123        ) = match clause {
124            MergeWhenPlan::UpdateMatched {
125                condition,
126                assignments,
127            } => (
128                condition.as_ref(),
129                assignments
130                    .iter()
131                    .flat_map(crate::plan::AssignmentPlan::expressions)
132                    .collect(),
133                &matched_schema,
134            ),
135            MergeWhenPlan::DeleteMatched { condition }
136            | MergeWhenPlan::NothingMatched { condition } => {
137                (condition.as_ref(), Vec::new(), &matched_schema)
138            }
139            MergeWhenPlan::UpdateNotMatchedBySource {
140                condition,
141                assignments,
142            } => (
143                condition.as_ref(),
144                assignments
145                    .iter()
146                    .flat_map(crate::plan::AssignmentPlan::expressions)
147                    .collect(),
148                target_schema,
149            ),
150            MergeWhenPlan::DeleteNotMatchedBySource { condition }
151            | MergeWhenPlan::NothingNotMatchedBySource { condition } => {
152                (condition.as_ref(), Vec::new(), target_schema)
153            }
154            MergeWhenPlan::InsertNotMatched {
155                condition,
156                columns,
157                values,
158                ..
159            } => (
160                condition.as_ref(),
161                columns
162                    .iter()
163                    .flat_map(crate::ast::AssignmentTarget::expressions)
164                    .chain(values)
165                    .collect(),
166                source_schema,
167            ),
168            MergeWhenPlan::NothingNotMatched { condition } => {
169                (condition.as_ref(), Vec::new(), source_schema)
170            }
171        };
172        if let Some(condition) = condition {
173            validate_boolean(condition, schema, "WHEN")?;
174        }
175        for expression in expressions {
176            expression_type(expression, schema)?;
177        }
178    }
179    Ok(())
180}
181
182pub fn expanded_merge_returning_projections(
183    catalog: &dyn super::returning::ReturningCatalog,
184    target_table: &str,
185    target_qualifier: &str,
186    aliases: &crate::ast::ReturningAliases,
187    source_schema: &crate::RowSchema,
188    source_relation: crate::ast::InternalRelationId,
189    returning: &[ProjectionPlan],
190) -> Result<Vec<ProjectionPlan>, SQLError> {
191    let target_star = ProjectionPlan {
192        expr: crate::ScalarExpr::QualifiedStar(target_qualifier.into()),
193        alias: None,
194    };
195    let target_projections = expanded_returning_projections(
196        catalog,
197        target_table,
198        target_qualifier,
199        aliases,
200        std::slice::from_ref(&target_star),
201    )?;
202    let mut projections = Vec::new();
203    for projection in returning {
204        match &projection.expr {
205            crate::ScalarExpr::Star => {
206                projections.extend(
207                    source_schema
208                        .columns()
209                        .iter()
210                        .enumerate()
211                        .filter(|(position, _)| {
212                            super::projection::visible_projection_source_position(
213                                source_schema,
214                                *position,
215                            )
216                        })
217                        .map(|(position, column)| ProjectionPlan {
218                            expr: crate::ScalarExpr::InternalColumn(
219                                source_relation.column(position),
220                            ),
221                            alias: Some(
222                                source_schema
223                                    .public_name(position)
224                                    .unwrap_or(column)
225                                    .to_string(),
226                            ),
227                        }),
228                );
229                projections.extend(target_projections.iter().cloned());
230            }
231            crate::ScalarExpr::QualifiedStar(qualifier)
232                if qualifier == target_qualifier
233                    || qualifier == &aliases.old
234                    || qualifier == &aliases.new =>
235            {
236                projections.extend(expanded_returning_projections(
237                    catalog,
238                    target_table,
239                    target_qualifier,
240                    aliases,
241                    std::slice::from_ref(projection),
242                )?);
243            }
244            _ => projections.push(projection.clone()),
245        }
246    }
247    Ok(projections)
248}
249
250pub fn merge_returning_source_schema(
251    source_schema: &crate::RowSchema,
252    source_relation: crate::ast::InternalRelationId,
253) -> crate::RowSchema {
254    let aliases = source_schema
255        .columns()
256        .iter()
257        .enumerate()
258        .map(|(position, _)| {
259            (
260                source_relation.column(position),
261                source_schema
262                    .physical_slot(position)
263                    .expect("source column has a physical slot"),
264                source_schema.column_type(position).cloned(),
265            )
266        })
267        .collect::<Vec<_>>();
268    crate::RowSchema::with_physical_internal_aliases(source_schema, &aliases)
269}
270
271/// The columns a `MERGE` supplies, as `PostgreSQL` records them in its target's `insertedCols` and `updatedCols`: those each `INSERT` action fills and those each `UPDATE` action sets.
272pub fn merge_target_columns(
273    catalog: &dyn super::mutation_privileges::MutationPrivilegeCatalog,
274    stmt: &MergePlan,
275) -> Result<Vec<String>, SQLError> {
276    let table_columns = catalog.bound_table_column_names(&stmt.target)?;
277    Ok(merge_supplied_columns(stmt, &table_columns))
278}
279
280/// The columns the actions of a `MERGE` supply, as [`merge_target_columns`] finds them, for a target whose columns `target_columns` lists in order.
281pub fn merge_supplied_columns(stmt: &MergePlan, target_columns: &[String]) -> Vec<String> {
282    let mut supplied = Vec::new();
283    for clause in &stmt.when_clauses {
284        match clause {
285            MergeWhenPlan::InsertNotMatched {
286                columns, values, ..
287            } => supplied.extend(merge_insert_columns(target_columns, columns, values)),
288            MergeWhenPlan::UpdateMatched { assignments, .. }
289            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => {
290                supplied.extend(
291                    assignments
292                        .iter()
293                        .flat_map(|assignment| assignment.target.column_names())
294                        .map(str::to_owned),
295                );
296            }
297            _ => {}
298        }
299    }
300    supplied
301}
302
303/// The columns an `INSERT` action fills: those it names, or else the leading columns of the table, one for each value; `DEFAULT VALUES` fills none.
304fn merge_insert_columns<T, V>(
305    table_columns: &[String],
306    columns: &[crate::ast::AssignmentTarget<T>],
307    values: &[V],
308) -> Vec<String> {
309    if columns.is_empty() {
310        table_columns.iter().take(values.len()).cloned().collect()
311    } else {
312        columns.iter().map(|target| target.column.clone()).collect()
313    }
314}
315
316pub fn ensure_merge_mutation_privileges(
317    catalog: &dyn super::mutation_privileges::MutationPrivilegeCatalog,
318    stmt: &MergePlan,
319) -> Result<(), SQLError> {
320    let mut column_privileges = BTreeSet::new();
321    let mut requires_delete = false;
322    let mut requires_any_insert = false;
323    let table_columns = catalog.bound_table_column_names(&stmt.target)?;
324    let privilege_subject = stmt
325        .target_privilege_subject
326        .clone()
327        .unwrap_or_else(|| catalog.current_role());
328    for clause in &stmt.when_clauses {
329        match clause {
330            MergeWhenPlan::InsertNotMatched {
331                columns, values, ..
332            } => {
333                if columns.is_empty() && values.is_empty() {
334                    requires_any_insert = true;
335                } else {
336                    let columns = merge_insert_columns(&table_columns, columns, values);
337                    column_privileges.extend(columns.into_iter().map(|column| {
338                        (
339                            crate::catalog::security::table::TableAclPrivilege::Insert,
340                            column,
341                        )
342                    }));
343                }
344            }
345            MergeWhenPlan::UpdateMatched { assignments, .. }
346            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => {
347                column_privileges.extend(
348                    assignments
349                        .iter()
350                        .flat_map(|assignment| assignment.target.targets())
351                        .map(|target| {
352                            (
353                                crate::catalog::security::table::TableAclPrivilege::Update,
354                                target.column.clone(),
355                            )
356                        }),
357                );
358            }
359            MergeWhenPlan::DeleteMatched { .. }
360            | MergeWhenPlan::DeleteNotMatchedBySource { .. } => requires_delete = true,
361            _ => {}
362        }
363    }
364    if requires_delete {
365        catalog.ensure_table_privilege_for(
366            &stmt.target,
367            &privilege_subject,
368            crate::catalog::security::table::TableAclPrivilege::Delete,
369        )?;
370    }
371    if requires_any_insert {
372        catalog.ensure_any_column_privilege_for(
373            &stmt.target,
374            &privilege_subject,
375            crate::catalog::security::table::TableAclPrivilege::Insert,
376        )?;
377    }
378    for (privilege, column) in column_privileges {
379        catalog.ensure_column_privilege_for(
380            &stmt.target,
381            &column,
382            &privilege_subject,
383            privilege,
384        )?;
385    }
386    Ok(())
387}
388
389pub fn merge_returning_schema(
390    routines: &dyn crate::routines::RoutineResolution,
391    catalog: &dyn super::returning::ReturningCatalog,
392    stmt: &MergePlan,
393    params: &[SQLParam],
394    source_schema: &crate::RowSchema,
395    ctes: &crate::binding::context::BindingContext<'_>,
396) -> Result<Option<crate::RowSchema>, SQLError> {
397    if stmt.returning.is_empty() {
398        return Ok(None);
399    }
400    let source_relation = crate::ast::InternalRelationId::allocate();
401    let projections = expanded_merge_returning_projections(
402        catalog,
403        &stmt.target,
404        &stmt.target_qualifier,
405        &stmt.returning_aliases,
406        source_schema,
407        source_relation,
408        &stmt.returning,
409    )?;
410    let returning_source_schema = merge_returning_source_schema(source_schema, source_relation);
411    let star_schema = super::returning::returning_target_schema(catalog, &stmt.target)?;
412    let expression_schema = super::returning_expression_schema(
413        &star_schema,
414        &stmt.target_qualifier,
415        &stmt.returning_aliases,
416        Some(&returning_source_schema),
417    );
418    crate::binding::analyze_projection_output_schema(
419        routines,
420        &projections,
421        &expression_schema,
422        &star_schema,
423        &stmt.subqueries,
424        params,
425        ctes,
426    )
427    .map(Some)
428}
429
430pub fn merge_command_returning_schema(
431    routines: &dyn crate::routines::RoutineResolution,
432    catalog: &dyn super::returning::ReturningCatalog,
433    rows: &dyn super::mutation_rows::MutationRowCatalog,
434    stmt: &MergePlan,
435    params: &[SQLParam],
436    bindings: &crate::binding::context::BindingContext<'_>,
437) -> Result<Option<crate::RowSchema>, SQLError> {
438    if stmt.returning.is_empty() {
439        return Ok(None);
440    }
441    let source =
442        crate::binding::analyze_source_plan_schema(routines, &stmt.source, params, bindings, None)?;
443    super::returning::validate_returning_alias_relations(
444        &stmt.target_qualifier,
445        &stmt.returning_aliases,
446        Some(&source),
447    )?;
448    let target =
449        super::mutation_rows::null_target_schema(rows, &stmt.target, &stmt.target_qualifier)?;
450    validate_merge_action_scopes(routines, stmt, &target, &source, params, bindings)?;
451    merge_returning_schema(routines, catalog, stmt, params, &source, bindings)
452}
453
454pub fn validate_merge_target_columns(
455    catalog: &dyn crate::assignment::columns::AssignmentColumnCatalog,
456    stmt: &MergePlan,
457) -> Result<(), SQLError> {
458    use crate::assignment::columns::validate_mutation_targets;
459    for clause in &stmt.when_clauses {
460        match clause {
461            MergeWhenPlan::UpdateMatched { assignments, .. }
462            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => {
463                validate_mutation_targets(
464                    catalog,
465                    &stmt.target,
466                    assignments
467                        .iter()
468                        .flat_map(|assignment| assignment.target.targets()),
469                    "MERGE UPDATE",
470                    false,
471                )?;
472            }
473            MergeWhenPlan::InsertNotMatched { columns, .. } => validate_mutation_targets(
474                catalog,
475                &stmt.target,
476                columns.iter(),
477                "MERGE INSERT",
478                true,
479            )?,
480            _ => {}
481        }
482    }
483    validate_merge_identity_targets(catalog, stmt)
484}
485
486/// Reject a `MERGE` update that assigns a `GENERATED ALWAYS` identity column anything but `DEFAULT`, and an insert that supplies a value for one without an `OVERRIDING` clause, whether or not a row reaches the action.
487fn validate_merge_identity_targets(
488    catalog: &dyn crate::assignment::columns::AssignmentColumnCatalog,
489    stmt: &MergePlan,
490) -> Result<(), SQLError> {
491    let identity = super::generated_values::GeneratedValueColumns::of(catalog, &stmt.target)?;
492    for clause in &stmt.when_clauses {
493        match clause {
494            MergeWhenPlan::UpdateMatched { assignments, .. }
495            | MergeWhenPlan::UpdateNotMatchedBySource { assignments, .. } => {
496                identity.validate_update(assignments.iter().flat_map(|assignment| {
497                    assignment.target.targets().iter().map(move |target| {
498                        (
499                            target.column.as_str(),
500                            matches!(assignment.value, crate::ScalarExpr::Default),
501                        )
502                    })
503                }))?;
504            }
505            MergeWhenPlan::InsertNotMatched {
506                columns,
507                overriding,
508                values,
509                ..
510            } => {
511                let targets = if columns.is_empty() {
512                    catalog
513                        .try_describe_table(&stmt.target)
514                        .map_err(|error| {
515                            SQLError::Internal(format!("read MERGE target columns: {error}"))
516                        })?
517                        .unwrap_or_default()
518                        .into_iter()
519                        .map(|column| column.name)
520                        .collect::<Vec<_>>()
521                } else {
522                    columns.iter().map(|target| target.column.clone()).collect()
523                };
524                identity.validate_insert(
525                    targets.iter().map(String::as_str).zip(
526                        values
527                            .iter()
528                            .map(|value| !matches!(value, crate::ScalarExpr::Default)),
529                    ),
530                    *overriding,
531                )?;
532            }
533            _ => {}
534        }
535    }
536    Ok(())
537}