Skip to main content

ruststream_sqlx_dialect/
postgres.rs

1//! The built-in Postgres dialect.
2
3use std::num::NonZeroUsize;
4
5use crate::advisory::Advisory;
6use crate::dialect::Dialect;
7use crate::form::KeyPart;
8use crate::insert::{Placeholder, Spelling, quote_into};
9use crate::lease::Lease;
10use crate::opening::{Isolation, Opening, Opens, level};
11use crate::outbox::OutboxDialect;
12use crate::row_lock::RowLock;
13use crate::spec::TableSpec;
14use crate::statement::{ClaimShape, NameLimit, Param, Statement, StatementError};
15use crate::table_name::TableName;
16use crate::writer::{BuiltIn, OutboxWriter, Probe, SqlWriter};
17
18/// How Postgres reads the current time: the start of the statement, so a settlement made long
19/// after the claim began records its own moment.
20const DATABASE_NOW: &str = "statement_timestamp()";
21
22/// How a claim locks the rows it takes: until its transaction ends, skipping the rows another
23/// claim holds.
24const LOCK: &str = " FOR UPDATE SKIP LOCKED";
25
26/// Tries the session's lock on the 64-bit hash of a key, without waiting.
27const TRY_LOCK: &str = "SELECT pg_try_advisory_lock(hashtextextended($1, 0))::int::bigint";
28
29/// Releases the session's lock on the 64-bit hash of a key.
30const UNLOCK: &str = "SELECT pg_advisory_unlock(hashtextextended($1, 0))::int::bigint";
31
32/// How Postgres spells an insert, for the `const fn` that writes it and for [`Postgres`]'s own
33/// statements.
34pub(crate) const SPELLING: Spelling = Spelling {
35    dialect: "postgres",
36    quote: b'"',
37    placeholder: Placeholder::Dollar,
38    default_row: " DEFAULT VALUES",
39    // `NAMEDATALEN` less its terminator. Postgres truncates a longer identifier without an
40    // error, so the statement would address another object.
41    name_limit: Some(NameLimit::Bytes(63)),
42};
43
44/// Postgres: double-quoted names, `$1` placeholders, rows claimed with `FOR UPDATE SKIP LOCKED`.
45///
46/// It builds the statements of the row lock form ([`RowLock`]), of the lease form ([`Lease`]), of
47/// the advisory lock form ([`Advisory`]), and the insert. A lease claim is one statement: it locks
48/// the claimable rows, writes their lease, and returns them as they were. Every name is quoted, so
49/// a name keeps its case and may hold any character; a name over 63 bytes, which Postgres would
50/// cut short without a word, is refused. A table on the database's clock reads
51/// `statement_timestamp()`.
52///
53/// In the advisory lock form a row's key is the text `concat` renders from the key's parts, a
54/// column without a value read as empty text. A session lock on the key's 64-bit hash
55/// (`hashtextextended`, Postgres 11 or later) holds the row, so two keys with one hash wait for
56/// each other: a delay, never a double delivery. The claim leaves out the keys another session
57/// holds: it probes each candidate's key with a shared lock that ends with the claim's own
58/// transaction, in claim order and only until it has its rows. The take counts the attempt and
59/// returns the row in one statement.
60///
61/// In a table with FIFO groups a claim's transaction takes its group
62/// ([`fifo_guard`](Dialect::fifo_guard)) with a lock the transaction holds on the 64-bit hash of
63/// the table's name and the group. Two groups with one hash wait for each other: a delay, never
64/// two rows of one group in work.
65///
66/// A table's transactions open with `BEGIN`, or at the isolation level it names
67/// ([`begin`](Dialect::begin)): READ COMMITTED, REPEATABLE READ or SERIALIZABLE. Postgres runs
68/// READ UNCOMMITTED as READ COMMITTED, so a table that names it is refused: a level the database
69/// does not keep is a level it lacks.
70///
71/// # Examples
72///
73/// A subscription asks Postgres for the statements of its table, as `ruststream-sqlx` does when the
74/// subscription starts:
75///
76/// ```
77/// use ruststream_sqlx_dialect::{
78///     ClaimShape, Column, Dialect, Form, Postgres, RowLock, Statement, StatementError, TableSpec,
79/// };
80///
81/// // The ledger as `#[derive(Inbox)]` describes it: one entry of an account in work at a time.
82/// const LEDGER: TableSpec<'static> =
83///     TableSpec::new("ledger", Column::new("entry_id"), Form::RowLock)
84///         .fifo_group(Column::new("account"));
85///
86/// // What a claim's transaction runs on a row lock table: the guard of its group, where the
87/// // table keeps groups in order, then the claim.
88/// fn claim_statements(
89///     dialect: &impl RowLock,
90///     spec: &TableSpec<'_>,
91/// ) -> Result<Vec<Statement>, StatementError> {
92///     let mut statements: Vec<Statement> = dialect.fifo_guard(spec)?.into_iter().collect();
93///     statements.push(dialect.lock_claim(spec, ClaimShape::Rows)?);
94///     Ok(statements)
95/// }
96///
97/// fn main() -> Result<(), StatementError> {
98///     // Postgres takes the account's group with a lock the transaction holds, then claims the
99///     // group's head.
100///     let statements = claim_statements(&Postgres, &LEDGER)?;
101///     assert_eq!(statements.len(), 2);
102///     Ok(())
103/// }
104/// ```
105#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
106pub struct Postgres;
107
108impl BuiltIn for Postgres {
109    const SPELLING: Spelling = SPELLING;
110
111    const ROW_LOCKS: bool = true;
112
113    /// `standard_conforming_strings`, on by default since Postgres 9.1, keeps a backslash as it is.
114    const BACKSLASH_ESCAPES: bool = false;
115
116    const UPDATE_RETURNS: bool = true;
117
118    /// A shared lock conflicts only with the session lock a delivery holds, so concurrent claims
119    /// probe one key without excluding each other.
120    const PROBE: Probe = Probe::Lock(
121        "pg_try_advisory_xact_lock_shared(hashtextextended(",
122        ", 0))",
123    );
124
125    /// Postgres widens a `smallint` that loses an `integer` one to `integer`, and a struct's
126    /// `i16` field reads no `integer`; one `smallint` less keeps every integer column's type.
127    const UNCOUNT: &'static str = " - 1::smallint";
128
129    fn database_now(&self) -> &'static str {
130        DATABASE_NOW
131    }
132
133    fn database_later(&self, sql: &mut SqlWriter<'_, Self>) {
134        // The cast types the delay where the statement is prepared. Prepared without parameter
135        // types, as a startup check prepares it, the server would read the delay as `double
136        // precision`, and an integer bound on that connection afterwards as almost nothing.
137        sql.push(DATABASE_NOW)
138            .push(" + ")
139            .param(Param::Delay)
140            .push("::bigint * interval '1 microsecond'");
141    }
142
143    // Postgres keeps its advisory locks per database, so the key alone names the lock.
144    fn render_lock_key(&self, sql: &mut SqlWriter<'_, Self>, _: Option<&str>, key: &[KeyPart<'_>]) {
145        sql.push("concat(")
146            .key_parts(key, ", ", |sql, column| {
147                sql.ident(column);
148            })
149            .push(")");
150    }
151}
152
153impl Dialect for Postgres {
154    fn name(&self) -> &'static str {
155        "postgres"
156    }
157
158    fn quote_into(&self, ident: &str, out: &mut String) {
159        quote_into(SPELLING.quote, ident, out);
160    }
161
162    fn placeholder_into(&self, index: NonZeroUsize, out: &mut String) {
163        out.push('$');
164        out.push_str(&index.to_string());
165    }
166
167    fn fetch(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
168        self.spec_fits(spec)?;
169        let id = spec.id().name();
170        let mut sql = SqlWriter::new(self);
171        sql.push("SELECT ")
172            .columns(spec)
173            .push(" FROM ")
174            .table(spec)
175            .push(" WHERE ")
176            .ident(id)
177            .push(" = ANY(")
178            .param(Param::Ids)
179            .push(")");
180        Ok(sql.finish())
181    }
182
183    fn ack(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
184        self.finish_statement(spec)
185    }
186
187    fn retry(&self, spec: &TableSpec<'_>) -> Result<Option<Statement>, StatementError> {
188        self.retry_statement(spec)
189    }
190
191    fn retry_after(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
192        self.retry_after_statement(spec)
193    }
194
195    fn discard(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
196        self.finish_statement(spec)
197    }
198
199    fn dead_letter_group(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
200        self.dead_letter_group_statement(spec)
201    }
202
203    fn dead_letter_table(
204        &self,
205        spec: &TableSpec<'_>,
206        target: TableName<'_>,
207    ) -> Result<Vec<Statement>, StatementError> {
208        self.movable(spec, target)?;
209        let mut sql = SqlWriter::new(self);
210        sql.push("WITH moved AS (DELETE FROM ")
211            .table(spec)
212            .settled_row(spec)
213            .push(" RETURNING ")
214            .columns(spec)
215            .push(") INSERT INTO ")
216            .table_name(target);
217        if !spec.selects_all() {
218            sql.push(" (").columns(spec).push(")");
219        }
220        sql.push(" SELECT ").moved_columns(spec).push(" FROM moved");
221        Ok(vec![sql.finish()])
222    }
223
224    fn insert(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
225        self.insert_statement(spec)
226    }
227
228    fn begin(&self, opening: Opening) -> Result<Option<&'static str>, StatementError> {
229        match opening {
230            Opening::Default => Ok(None),
231            Opening::Isolation(Isolation::ReadCommitted) => {
232                Ok(Some("BEGIN ISOLATION LEVEL READ COMMITTED"))
233            }
234            Opening::Isolation(Isolation::RepeatableRead) => {
235                Ok(Some("BEGIN ISOLATION LEVEL REPEATABLE READ"))
236            }
237            Opening::Isolation(Isolation::Serializable) => {
238                Ok(Some("BEGIN ISOLATION LEVEL SERIALIZABLE"))
239            }
240            // Postgres runs READ UNCOMMITTED as READ COMMITTED: a declared level it does not keep
241            // is a level it lacks.
242            Opening::Isolation(Isolation::ReadUncommitted) | Opening::Mode(_) => {
243                Err(opening.refused(self.name()))
244            }
245        }
246    }
247
248    fn fifo_guard(&self, spec: &TableSpec<'_>) -> Result<Option<Statement>, StatementError> {
249        if !self.guards_group(spec)? {
250            return Ok(None);
251        }
252        // The group's name: the table as the statements name it, unquoted, then `:` and the
253        // group, so the groups of two tables are two names.
254        let mut table = spec
255            .schema()
256            .map(|schema| format!("{schema}."))
257            .unwrap_or_default();
258        table.push_str(spec.table());
259        table.push(':');
260        let mut sql = SqlWriter::new(self);
261        sql.push("SELECT pg_try_advisory_xact_lock(hashtextextended(")
262            .literal(&table)
263            .push(" || ")
264            .param(Param::Group)
265            .push(", 0))::int::bigint");
266        Ok(Some(sql.finish()))
267    }
268}
269
270impl OutboxDialect for Postgres {
271    fn outbox_fetch(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
272        self.outbox_fetch_statement(spec)
273    }
274
275    fn outbox_mark(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
276        self.outbox_mark_statement(spec)
277    }
278
279    fn outbox_recover(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
280        self.outbox_recover_statement(spec)
281    }
282}
283
284impl Opens<level::ReadCommitted> for Postgres {}
285
286impl Opens<level::RepeatableRead> for Postgres {}
287
288impl Opens<level::Serializable> for Postgres {}
289
290impl RowLock for Postgres {
291    fn lock_claim(
292        &self,
293        spec: &TableSpec<'_>,
294        shape: ClaimShape,
295    ) -> Result<Statement, StatementError> {
296        self.locked(spec, "lock_claim")?;
297        let mut sql = SqlWriter::new(self);
298        sql.claim(spec, shape, LOCK);
299        Ok(sql.finish())
300    }
301}
302
303impl Advisory for Postgres {
304    fn advisory_claim(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
305        self.advisory_claim_statement(spec)
306    }
307
308    fn lock(&self) -> Option<Statement> {
309        Some(Statement::new(TRY_LOCK, [Param::Key]))
310    }
311
312    fn unlock(&self) -> Option<Statement> {
313        Some(Statement::new(UNLOCK, [Param::Key]))
314    }
315
316    fn take(
317        &self,
318        spec: &TableSpec<'_>,
319        shape: ClaimShape,
320    ) -> Result<Vec<Statement>, StatementError> {
321        self.take_statements(spec, shape)
322    }
323}
324
325impl Lease for Postgres {
326    fn lease_claim(
327        &self,
328        spec: &TableSpec<'_>,
329        shape: ClaimShape,
330    ) -> Result<Statement, StatementError> {
331        let expiry = self.leased(spec, "lease_claim")?;
332        let mut sql = SqlWriter::new(self);
333        sql.lease_claim(spec, shape, expiry.name(), LOCK);
334        Ok(sql.finish())
335    }
336
337    fn extend(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
338        self.extend_statement(spec)
339    }
340
341    fn stamp(&self, spec: &TableSpec<'_>) -> Result<Statement, StatementError> {
342        self.stamp_statement(spec)
343    }
344}