Skip to main content

uqa_sql/semantics/
locking.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! SQL row-lock restrictions, propagation, and mutation lock strength.
8
9use crate::ast::{ColumnDef, LockStrength, LockWait, LockingClause, TableKeyConstraint};
10use crate::{
11    plan::{
12        locking::ResolvedRowLock, ComputePlan, QueryBlockPlan, QueryPlan, RelationalPlan,
13        SourcePlan,
14    },
15    SQLError,
16};
17use std::collections::BTreeSet;
18
19pub mod null_rejection;
20
21fn cte_plan_has_row_locks(body: &crate::plan::CtePlanBody) -> bool {
22    match body {
23        crate::plan::CtePlanBody::Query(query) => query_plan_has_row_locks(query),
24        crate::plan::CtePlanBody::Command(command) => {
25            command
26                .ctes()
27                .iter()
28                .any(|cte| cte_plan_has_row_locks(&cte.body))
29                || command
30                    .query_inputs()
31                    .into_iter()
32                    .any(query_plan_has_row_locks)
33                || command
34                    .source_input()
35                    .is_some_and(source_plan_has_row_locks)
36        }
37    }
38}
39
40pub fn query_plan_has_row_locks(query: &QueryPlan) -> bool {
41    query
42        .ctes
43        .iter()
44        .any(|cte| cte_plan_has_row_locks(&cte.body))
45        || relational_has_row_locks(&query.root)
46}
47
48fn relational_has_row_locks(plan: &RelationalPlan) -> bool {
49    match plan {
50        RelationalPlan::QueryBlock(block) => {
51            !block.locking.is_empty()
52                || block.from.as_ref().is_some_and(source_plan_has_row_locks)
53                || block.subqueries.iter().any(query_plan_has_row_locks)
54        }
55        RelationalPlan::SetOp { left, right, .. } => {
56            query_plan_has_row_locks(left) || query_plan_has_row_locks(right)
57        }
58        RelationalPlan::Values { .. } => false,
59    }
60}
61
62fn source_plan_has_row_locks(source: &SourcePlan) -> bool {
63    match source {
64        SourcePlan::Join { left, right, .. } => {
65            source_plan_has_row_locks(left) || source_plan_has_row_locks(right)
66        }
67        SourcePlan::Subquery { body, .. } => query_plan_has_row_locks(body),
68        SourcePlan::Table { .. }
69        | SourcePlan::Values { .. }
70        | SourcePlan::Function { .. }
71        | SourcePlan::FunctionGroup { .. } => false,
72    }
73}
74
75/// Apply a row mark selected for a stored view to the view plan before execution. Stored view plans are not present when the SQL compiler pushes row marks into derived tables, so runtime expansion must perform the same propagation to ensure an outer `NOWAIT` or `SKIP LOCKED` policy is merged before an inner row mark can block.
76pub fn apply_propagated_view_lock(plan: &mut QueryPlan, target: &ResolvedRowLock) {
77    apply_propagated_lock_to_relational(&mut plan.root, target.strength, target.wait);
78}
79
80fn apply_propagated_lock_to_relational(
81    plan: &mut RelationalPlan,
82    strength: LockStrength,
83    wait: LockWait,
84) {
85    let RelationalPlan::QueryBlock(block) = plan else {
86        return;
87    };
88    block.locking.push(LockingClause {
89        strength,
90        wait,
91        relations: Vec::new(),
92    });
93    if let Some(source) = block.from.as_mut() {
94        apply_propagated_lock_to_subqueries(source, strength, wait);
95    }
96}
97
98fn apply_propagated_lock_to_subqueries(
99    source: &mut SourcePlan,
100    strength: LockStrength,
101    wait: LockWait,
102) {
103    match source {
104        SourcePlan::Join { left, right, .. } => {
105            apply_propagated_lock_to_subqueries(left, strength, wait);
106            apply_propagated_lock_to_subqueries(right, strength, wait);
107        }
108        SourcePlan::Subquery { body, .. } => {
109            apply_propagated_lock_to_relational(&mut body.root, strength, wait);
110        }
111        SourcePlan::Table { .. }
112        | SourcePlan::Values { .. }
113        | SourcePlan::Function { .. }
114        | SourcePlan::FunctionGroup { .. } => {}
115    }
116}
117
118pub fn validate_locking_block_shape(
119    block: &QueryBlockPlan,
120    strength: LockStrength,
121) -> Result<(), SQLError> {
122    let label = strength.sql_name();
123    if block.distinct || !block.distinct_on.is_empty() {
124        return Err(SQLError::Unsupported(format!(
125            "{label} is not allowed with DISTINCT clause"
126        )));
127    }
128    if !block.group_by.is_empty() || !block.grouping_sets.is_empty() {
129        return Err(SQLError::Unsupported(format!(
130            "{label} is not allowed with GROUP BY clause"
131        )));
132    }
133    if block.having.is_some() {
134        return Err(SQLError::Unsupported(format!(
135            "{label} is not allowed with HAVING clause"
136        )));
137    }
138    if matches!(block.compute, ComputePlan::Window)
139        || block
140            .order_by
141            .iter()
142            .any(|ordering| ordering.expr.contains_window())
143    {
144        return Err(SQLError::Unsupported(format!(
145            "{label} is not allowed with window functions"
146        )));
147    }
148    if matches!(block.compute, ComputePlan::Aggregate) {
149        return Err(SQLError::Unsupported(format!(
150            "{label} is not allowed with aggregate functions"
151        )));
152    }
153    Ok(())
154}
155
156pub fn update_lock_strength(
157    keys: &[TableKeyConstraint],
158    definitions: &[ColumnDef],
159    columns: &[String],
160) -> LockStrength {
161    let assigned = columns.iter().map(String::as_str).collect::<BTreeSet<_>>();
162    let touches_key = keys.iter().any(|constraint| {
163        constraint.columns.iter().any(|column| {
164            if assigned.contains(column.as_str()) {
165                return true;
166            }
167            let Some(generated) = definitions
168                .iter()
169                .find(|definition| definition.name == *column)
170                .and_then(|definition| definition.generated.as_ref())
171            else {
172                return false;
173            };
174            let mut dependencies = BTreeSet::new();
175            let expression = crate::plan::ExpressionPlan::lower((*generated.expression).clone());
176            !expression.scalar.collect_columns(&mut dependencies)
177                || dependencies
178                    .iter()
179                    .any(|dependency| assigned.contains(dependency.as_str()))
180        })
181    });
182    if touches_key {
183        crate::ast::LockStrength::ForUpdate
184    } else {
185        crate::ast::LockStrength::ForNoKeyUpdate
186    }
187}