uqa_sql/semantics/
locking.rs1use 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
75pub 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}