1use rupg_common::{Error, Result, SqlState};
6use rupg_sql::nodes::{LockClauseStrength, LockingClause, Node};
7
8use crate::coerce::AtOpt;
9use crate::from::{FromItem, Join, Relation};
10use crate::setop::strength_name;
11use crate::typename::place;
12use crate::{JoinKind, Query};
13
14#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct RowMark {
17 pub relation: usize,
19 pub strength: LockClauseStrength,
20}
21
22enum Entry {
24 Relation(usize),
25 Join(Option<String>),
26}
27
28fn entries(query: &Query) -> Vec<Entry> {
30 fn walk(item: &FromItem, out: &mut Vec<Entry>) {
31 match item {
32 FromItem::Relation(index) => out.push(Entry::Relation(*index)),
33 FromItem::Join(join) => {
34 walk(&join.left, out);
35 walk(&join.right, out);
36 out.push(Entry::Join(join_name(join)));
37 }
38 }
39 }
40 let mut out = Vec::new();
41 for item in &query.from {
42 walk(item, &mut out);
43 }
44 out
45}
46
47fn join_name(join: &Join) -> Option<String> {
49 join.alias.clone().or_else(|| join.using_alias.clone())
50}
51
52fn relation_entries(query: &Query) -> Vec<usize> {
54 entries(query)
55 .into_iter()
56 .filter_map(|entry| match entry {
57 Entry::Relation(index) => Some(index),
58 Entry::Join(_) => None,
59 })
60 .collect()
61}
62
63fn is_relation(relation: &Relation) -> bool {
65 relation.oid != 0 && relation.function.is_none() && relation.cte.is_none()
66}
67
68fn is_subquery(relation: &Relation) -> bool {
70 relation.oid == 0 && relation.subquery.is_some() && relation.cte.is_none()
71}
72
73pub(crate) fn check_select_locking(query: &Query, strength: LockClauseStrength) -> Result<()> {
75 let what = if query.set_op.is_some() {
76 "UNION/INTERSECT/EXCEPT"
77 } else if !query.distinct.is_empty() {
78 "DISTINCT clause"
79 } else if !query.group.is_empty()
80 || (query.grouped && !has_aggs(query) && query.having.is_none())
81 {
82 "GROUP BY clause"
83 } else if query.having.is_some() {
84 "HAVING clause"
85 } else if has_aggs(query) {
86 "aggregate functions"
87 } else if query.targets.iter().any(|t| t.expr.first_window().is_some()) {
88 "window functions"
89 } else if query.has_target_srfs() {
90 "set-returning functions in the target list"
91 } else {
92 return Ok(());
93 };
94 Err(Error::new(
95 SqlState::FEATURE_NOT_SUPPORTED,
96 format!("{} is not allowed with {what}", strength_name(strength)),
97 ))
98}
99
100fn has_aggs(query: &Query) -> bool {
102 query.targets.iter().any(|t| t.expr.first_agg().is_some())
103 || query.having.as_ref().is_some_and(|h| h.first_agg().is_some())
104}
105
106fn apply(query: &mut Query, relation: usize, strength: LockClauseStrength) {
108 match query.row_marks.iter_mut().find(|mark| mark.relation == relation) {
109 Some(mark) => {
110 if strength.0 > mark.strength.0 {
111 mark.strength = strength;
112 }
113 }
114 None => query.row_marks.push(RowMark { relation, strength }),
115 }
116}
117
118fn apply_relation(query: &mut Query, relation: usize, strength: LockClauseStrength) -> Result<()> {
120 apply(query, relation, strength);
121 if let Some(view) = query.relations[relation].subquery.as_deref_mut() {
122 mark_for_locking(view, strength)?;
123 }
124 Ok(())
125}
126
127fn mark_for_locking(query: &mut Query, strength: LockClauseStrength) -> Result<()> {
129 for index in relation_entries(query) {
130 let relation = &query.relations[index];
131 if is_relation(relation) {
132 apply_relation(query, index, strength)?;
133 } else if is_subquery(relation) {
134 apply(query, index, strength);
135 if let Some(sub) = query.relations[index].subquery.as_deref_mut() {
136 mark_for_locking(sub, strength)?;
137 }
138 }
139 }
140 if let Some(first) = query.row_marks.first() {
141 check_select_locking(query, first.strength)?;
142 }
143 Ok(())
144}
145
146pub(crate) fn transform_locking_clause(query: &mut Query, clause: &LockingClause) -> Result<()> {
148 let strength = clause.strength;
149 check_select_locking(query, strength)?;
150 if clause.lockedRels.is_empty() {
151 for index in relation_entries(query) {
152 lock_relation(query, index, strength)?;
153 }
154 return Ok(());
155 }
156 let order = entries(query);
157 for node in clause.lockedRels.iter().flatten() {
158 let Node::RangeVar(rv) = node else { continue };
159 let at = place(rv.location);
160 if rv.catalogname.is_some() || rv.schemaname.is_some() {
161 return Err(Error::new(
162 SqlState::SYNTAX_ERROR,
163 format!("{} must specify unqualified relation names", strength_name(strength)),
164 )
165 .at_opt(at));
166 }
167 let name = rv.relname.as_deref().unwrap_or_default();
168 let found = order.iter().find_map(|entry| match entry {
169 Entry::Join(alias) => {
170 (alias.as_deref() == Some(name)).then_some(Found::Error("a join"))
171 }
172 Entry::Relation(index) => {
173 let relation = &query.relations[*index];
174 let hidden = relation.alias.is_none()
175 && (is_subquery(relation) || relation.values.is_some());
176 if hidden || relation.name != name {
177 return None;
178 }
179 Some(if relation.cte.is_some() {
180 Found::Error("a WITH query")
181 } else if relation.json_table().is_some() {
182 Found::Error("a table function")
183 } else if relation.function.is_some() {
184 Found::Error("a function")
185 } else if relation.values.is_some() {
186 Found::Error("VALUES")
187 } else {
188 Found::Relation(*index)
189 })
190 }
191 });
192 match found {
193 Some(Found::Relation(index)) => lock_relation(query, index, strength)?,
194 Some(Found::Error(what)) => {
195 return Err(Error::new(
196 SqlState::FEATURE_NOT_SUPPORTED,
197 format!("{} cannot be applied to {what}", strength_name(strength)),
198 )
199 .at_opt(at));
200 }
201 None => {
202 return Err(Error::new(
203 SqlState::UNDEFINED_TABLE,
204 format!(
205 "relation \"{name}\" in {} clause not found in FROM clause",
206 strength_name(strength)
207 ),
208 )
209 .at_opt(at));
210 }
211 }
212 }
213 Ok(())
214}
215
216enum Found {
218 Relation(usize),
219 Error(&'static str),
220}
221
222fn lock_relation(query: &mut Query, index: usize, strength: LockClauseStrength) -> Result<()> {
224 let relation = &query.relations[index];
225 if is_relation(relation) {
226 apply_relation(query, index, strength)
227 } else if is_subquery(relation) {
228 apply(query, index, strength);
229 let all = LockingClause { lockedRels: Vec::new(), strength, ..LockingClause::default() };
230 match query.relations[index].subquery.as_deref_mut() {
231 Some(sub) => transform_locking_clause(sub, &all),
232 None => Ok(()),
233 }
234 } else {
235 Ok(())
236 }
237}
238
239pub fn check_outer_joins(query: &Query) -> Result<()> {
245 fn relations(item: &FromItem, out: &mut Vec<usize>) {
246 match item {
247 FromItem::Relation(index) => out.push(*index),
248 FromItem::Join(join) => {
249 relations(&join.left, out);
250 relations(&join.right, out);
251 }
252 }
253 }
254 fn check(query: &Query, item: &FromItem) -> Result<()> {
255 let FromItem::Join(join) = item else { return Ok(()) };
256 check(query, &join.left)?;
257 check(query, &join.right)?;
258 let (mut left, mut right) = (Vec::new(), Vec::new());
259 relations(&join.left, &mut left);
260 relations(&join.right, &mut right);
261 let nullable: Vec<usize> = match join.kind {
262 JoinKind::Left => right,
263 JoinKind::Right => left,
264 JoinKind::Full => left.into_iter().chain(right).collect(),
265 _ => Vec::new(),
266 };
267 match query.row_marks.iter().find(|mark| nullable.contains(&mark.relation)) {
268 Some(mark) => Err(Error::new(
269 SqlState::FEATURE_NOT_SUPPORTED,
270 format!(
271 "{} cannot be applied to the nullable side of an outer join",
272 strength_name(mark.strength)
273 ),
274 )),
275 None => Ok(()),
276 }
277 }
278 for item in &query.from {
279 check(query, item)?;
280 }
281 for relation in &query.relations {
282 if let Some(sub) = &relation.subquery {
283 check_outer_joins(sub)?;
284 }
285 }
286 Ok(())
287}
288
289pub fn first_row_mark(query: &Query) -> Option<LockClauseStrength> {
291 let own = query.row_marks.iter().find(|mark| query.relations[mark.relation].subquery.is_none());
292 match own {
293 Some(mark) => Some(mark.strength),
294 None => {
295 query.relations.iter().filter_map(|r| r.subquery.as_deref()).find_map(first_row_mark)
296 }
297 }
298}
299
300pub fn locked_tables(query: &Query, out: &mut Vec<u32>) {
302 for mark in &query.row_marks {
303 let relation = &query.relations[mark.relation];
304 if relation.oid != 0 && relation.subquery.is_none() {
305 out.push(relation.oid);
306 }
307 }
308 for sub in query.relations.iter().filter_map(|r| r.subquery.as_deref()) {
309 locked_tables(sub, out);
310 }
311}