Skip to main content

rupg_analyze/
locking.rs

1//! `FOR UPDATE`, `FOR NO KEY UPDATE`, `FOR SHARE` and `FOR KEY SHARE`, as `transformLockingClause` and `CheckSelectLocking` of `analyze.c` and `markQueryForLocking` of `rewriteHandler.c` run them.
2//!
3//! rupg takes no lock on a row. No statement of rupg changes the rows of a table, so a lock on a row has no effect on another session. The analyzer gives the errors of PostgreSQL and keeps the row marks, which the planner and the executor check.
4
5use 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/// A relation that `FOR UPDATE` or `FOR SHARE` locks, as a `RowMarkClause`.
15#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct RowMark {
17    /// The index of the relation in [`Query::relations`].
18    pub relation: usize,
19    pub strength: LockClauseStrength,
20}
21
22/// An item of the range table of a query in the order of PostgreSQL: a relation of [`Query::relations`], or a join with the name that it has as a relation, which comes after its two sides.
23enum Entry {
24    Relation(usize),
25    Join(Option<String>),
26}
27
28/// The items of the range table of the query.
29fn 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
47/// The name of a join as a relation: its alias, or the alias of `USING (...) AS name`. A join with neither has no name.
48fn join_name(join: &Join) -> Option<String> {
49    join.alias.clone().or_else(|| join.using_alias.clone())
50}
51
52/// The relations of the range table of the query, in its order.
53fn 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
63/// A table or a view: an `RTE_RELATION` before the rewriter reads the view.
64fn is_relation(relation: &Relation) -> bool {
65    relation.oid != 0 && relation.function.is_none() && relation.cte.is_none()
66}
67
68/// A subquery in `FROM`, which is not a view.
69fn is_subquery(relation: &Relation) -> bool {
70    relation.oid == 0 && relation.subquery.is_some() && relation.cte.is_none()
71}
72
73/// `CheckSelectLocking`: the parts of a query that cannot go with a lock on its rows.
74pub(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
100/// `hasAggs`: the query has an aggregate call of its own level.
101fn 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
106/// `applyLockingClause`: the row mark of a relation. A relation with two marks keeps the strongest one.
107fn 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
118/// The row mark of a table or a view. The rewriter puts the lock of a view on the relations of its query, as `markQueryForLocking` does, and the planner then checks the query of the view.
119fn 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
127/// `markQueryForLocking` on the query of a view, then the check of `preprocess_rowmarks` when the query has a row mark.
128fn 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
146/// `transformLockingClause`: the row marks of a locking clause of a query, on the relations that the clause names or on all the relations of `FROM`.
147pub(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
216/// The relation that a name of a locking clause finds: a table, a view or a subquery, or the error of another kind of relation.
217enum Found {
218    Relation(usize),
219    Error(&'static str),
220}
221
222/// The lock of a relation of `FROM`: the row mark of a table or a view, or the row mark of a subquery and the lock of all the relations of the subquery. The other relations have no lock.
223fn 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
239/// The check of `make_outerjoininfo`: a relation with a row mark cannot be on the nullable side of an outer join, at each level of the query.
240///
241/// # Errors
242///
243/// `0A000` for a row mark on the nullable side of an outer join.
244pub 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
289/// The strength of the first row mark on a table in the query, in its subqueries in `FROM` and in its views, as the first `PlanRowMark` of the plan.
290pub 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
300/// The tables with a row mark in the query, in its subqueries in `FROM` and in its views.
301pub 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}