Line data Source code
1 : // SPDX-License-Identifier: EUPL-1.2
2 : //! Postgres foundation: ONE shared pool, embedded migrations,
3 : //! transaction-scoped tenancy. Store implementations build on top.
4 :
5 : pub mod doc;
6 : pub mod entity;
7 : pub mod maintenance;
8 : pub mod outbox;
9 : pub mod temporal;
10 :
11 : use antares_model::TenantId;
12 : use serde_json::Value;
13 : use sqlx::postgres::{PgPool, PgPoolOptions};
14 : use sqlx::{Postgres, Transaction};
15 :
16 : /// Embedded migrations, run at start (like Scorpio's Flyway, but once).
17 : ///
18 : /// `ANTARES_MIGRATE=0`/`false` skips that run on this process, so serving
19 : /// replicas do not race the same DDL and a deployment can migrate once from a
20 : /// separate job or init container against the same database. Unset — the
21 : /// default — migrates on boot exactly as before.
22 : pub static MIGRATOR: sqlx::migrate::Migrator = sqlx::migrate!("./migrations");
23 :
24 : /// Is this connect failure a database whose migration history can never match
25 : /// this binary? A database still booting is worth waiting for; a history from
26 : /// another release is not, and a caller that retries it reports the timeout
27 : /// instead of the cause. Only the history errors are permanent: an execution
28 : /// failure can be a lock the next attempt gets.
29 18 : pub fn is_schema_mismatch(e: &sqlx::Error) -> bool {
30 : use sqlx::migrate::MigrateError::{
31 : Dirty, VersionMismatch, VersionMissing, VersionNotPresent, VersionTooNew, VersionTooOld,
32 : };
33 18 : match e {
34 14 : sqlx::Error::Migrate(m) => matches!(
35 14 : **m,
36 : VersionMissing(_)
37 : | VersionMismatch(_)
38 : | VersionNotPresent(_)
39 : | VersionTooOld(..)
40 : | VersionTooNew(..)
41 : | Dirty(_)
42 : ),
43 4 : _ => false,
44 : }
45 18 : }
46 :
47 : /// The `ANTARES_MIGRATE` switch: off only for an explicit `0`/`false`.
48 161 : fn migrate_enabled(v: Option<&str>) -> bool {
49 161 : !matches!(v, Some("0" | "false"))
50 161 : }
51 :
52 : /// One shared pool for all tenants — never per-tenant pools.
53 : /// `max_connections` ≈ 2× the PG box's cores; the default suits a small dev
54 : /// Postgres, deployments size it via config.
55 : ///
56 : /// Every acquisition and session is bounded: a saturated pool fails the one
57 : /// request after 5 s instead of queueing forever; idle/aged connections are
58 : /// recycled; and each session carries `statement_timeout`/`lock_timeout` so
59 : /// a runaway query or lost lock can never wedge a pooled connection.
60 154 : pub async fn connect(url: &str, max_connections: u32) -> Result<PgPool, sqlx::Error> {
61 154 : connect_with(url, max_connections, std::time::Duration::from_secs(30)).await
62 154 : }
63 :
64 : /// [`connect`] with the per-session `statement_timeout` chosen by the
65 : /// deployment (`ANTARES_PG_STATEMENT_TIMEOUT_MS`).
66 157 : pub async fn connect_with(
67 157 : url: &str,
68 157 : max_connections: u32,
69 157 : statement_timeout: std::time::Duration,
70 157 : ) -> Result<PgPool, sqlx::Error> {
71 : use std::time::Duration;
72 157 : let session = format!(
73 : "SET statement_timeout = '{}ms'; SET lock_timeout = '5s'",
74 157 : statement_timeout.as_millis()
75 : );
76 157 : let pool = PgPoolOptions::new()
77 157 : .max_connections(max_connections)
78 157 : .acquire_timeout(Duration::from_secs(5))
79 157 : .idle_timeout(Duration::from_secs(600))
80 157 : .max_lifetime(Duration::from_secs(1800))
81 476 : .after_connect(move |conn, _meta| {
82 470 : let session = session.clone();
83 470 : Box::pin(async move {
84 : use sqlx::Executor;
85 : // the only interpolated value is a validated integer of
86 : // milliseconds (parsed by the broker), never client text
87 470 : conn.execute(sqlx::raw_sql(sqlx::AssertSqlSafe(session)))
88 470 : .await?;
89 470 : Ok(())
90 470 : })
91 470 : })
92 157 : .connect(url)
93 157 : .await?;
94 : // DDL is exempt from the per-session statement timeout: building an index
95 : // on a large attr_instances legitimately runs longer than a query ever
96 : // should. The lock timeout stays — a migration blocked on a lock must
97 : // fail rather than hold the boot path.
98 : //
99 : // `SET` is session-scoped and the pool has no release hook, so the
100 : // connection is DETACHED first: it never goes back into the pool, and no
101 : // request is ever served by a connection with the runaway-query wall
102 : // switched off. The pool opens a fresh one (running `after_connect`) on
103 : // demand.
104 151 : if migrate_enabled(std::env::var("ANTARES_MIGRATE").ok().as_deref()) {
105 : use sqlx::{Connection, Executor};
106 151 : let mut migrate_conn = pool.acquire().await?.detach();
107 151 : migrate_conn
108 151 : .execute(sqlx::raw_sql("SET statement_timeout = 0"))
109 151 : .await?;
110 151 : MIGRATOR.run(&mut migrate_conn).await?;
111 151 : migrate_conn.close().await?;
112 0 : }
113 151 : Ok(pool)
114 157 : }
115 :
116 : /// Make RLS effective for this transaction — SET LOCAL only (transaction
117 : /// scoped), so a recycled pooled connection carries no tenant residue.
118 : /// Call first in EVERY transaction that touches tenant data.
119 78402 : pub async fn set_tenant(
120 78402 : tx: &mut Transaction<'_, Postgres>,
121 78402 : tenant: &TenantId,
122 78402 : ) -> Result<(), sqlx::Error> {
123 78402 : sqlx::query(crate::SET_TENANT_SQL)
124 78402 : .bind(tenant.as_str())
125 78402 : .execute(&mut **tx)
126 78402 : .await
127 78402 : .map(|_| ())
128 78402 : }
129 :
130 : /// Does the connected role bypass RLS (superuser or BYPASSRLS)? RLS is
131 : /// a belt only when the role wears it — the broker warns at startup when the
132 : /// belt is off.
133 : ///
134 : /// A probe that cannot answer fails CLOSED: an unreachable or erroring
135 : /// database is reported as bypassing, so the strict gate refuses the boot
136 : /// instead of passing on an error.
137 2 : pub async fn role_bypasses_rls(pool: &PgPool) -> bool {
138 2 : bypasses(
139 2 : sqlx::query_scalar(
140 2 : "SELECT rolsuper OR rolbypassrls FROM pg_roles WHERE rolname = current_user",
141 2 : )
142 2 : .fetch_one(pool)
143 2 : .await,
144 : )
145 2 : }
146 :
147 : /// What the connected server is, read ONCE at startup and served from
148 : /// `/q/health` afterwards: server version plus the version of each extension
149 : /// the broker's behaviour depends on. A probe that fails is not fatal — the
150 : /// broker serves fine without knowing its own server version — so the caller
151 : /// gets an empty object and health simply says nothing.
152 4 : pub async fn version_info(pool: &PgPool) -> Value {
153 : // The pool's own shape is what an operator sizes against, and it is
154 : // known without asking the server: how many connections this process may
155 : // hold, and how long a request waits for one before the broker answers
156 : // 503. Both are reported even when the version probe fails.
157 4 : let mut m = serde_json::Map::new();
158 4 : m.insert("engine".into(), "postgres".into());
159 4 : m.insert(
160 4 : "poolSize".into(),
161 4 : pool.options().get_max_connections().into(),
162 : );
163 4 : m.insert(
164 4 : "poolAcquireTimeoutSeconds".into(),
165 4 : pool.options().get_acquire_timeout().as_secs().into(),
166 : );
167 4 : let row: Result<(String, Option<String>, Option<String>), sqlx::Error> = sqlx::query_as(
168 4 : "SELECT current_setting('server_version'), \
169 4 : (SELECT extversion FROM pg_extension WHERE extname = 'postgis'), \
170 4 : (SELECT extversion FROM pg_extension WHERE extname = 'timescaledb')",
171 4 : )
172 4 : .fetch_one(pool)
173 4 : .await;
174 4 : match row {
175 4 : Ok((server, postgis, timescale)) => {
176 4 : m.insert("server".into(), server.into());
177 4 : if let Some(v) = postgis {
178 4 : m.insert("postgis".into(), v.into());
179 4 : }
180 4 : if let Some(v) = timescale {
181 2 : m.insert("timescaledb".into(), v.into());
182 4 : }
183 : }
184 0 : Err(e) => {
185 0 : tracing::warn!("server version probe failed ({e}); /q/health will not report it");
186 : }
187 : }
188 4 : Value::Object(m)
189 4 : }
190 :
191 : /// Open a transaction, timing what it cost to get one. `Pool::begin` is the
192 : /// only way to a `Transaction<'static>` (sqlx owns the pooled connection
193 : /// inside it), so the measurement covers the pool wait plus one BEGIN round
194 : /// trip; the round trip is sub-millisecond and the wait is what grows, up to
195 : /// the acquire timeout. Every transactional store call goes through here, so
196 : /// `antares_pg_transaction_begin_seconds` is where pool pressure shows.
197 77582 : pub(crate) async fn begin(
198 77582 : pool: &PgPool,
199 77582 : ) -> Result<sqlx::Transaction<'static, sqlx::Postgres>, sqlx::Error> {
200 77582 : let t0 = std::time::Instant::now();
201 77582 : let tx = pool.begin().await;
202 77582 : metrics::histogram!("antares_pg_transaction_begin_seconds").record(t0.elapsed().as_secs_f64());
203 77582 : tx
204 77582 : }
205 :
206 10 : fn bypasses(probe: Result<bool, sqlx::Error>) -> bool {
207 10 : probe.unwrap_or_else(|e| {
208 4 : tracing::error!(
209 : "row-level-security probe failed ({e}) — treating the role as RLS-bypassing"
210 : );
211 4 : true
212 4 : })
213 10 : }
214 :
215 : /// Arm the transaction-scoped `antares.service` escape (0001_init.sql) for
216 : /// the two internal cross-tenant jobs: outbox drain and temporal retention.
217 : /// NEVER call from a request path — request queries carry explicit tenant
218 : /// predicates and run under `set_tenant` only.
219 2366 : pub async fn set_service(tx: &mut Transaction<'_, Postgres>) -> Result<(), sqlx::Error> {
220 2366 : sqlx::query("SELECT set_config('antares.service', 'on', true)")
221 2366 : .execute(&mut **tx)
222 2366 : .await
223 2366 : .map(|_| ())
224 2366 : }
225 :
226 : /// The write's own claim on its tenant (ADR-0001: "every implicit tenant
227 : /// creation inserts its row in the same transaction as the document"). Two
228 : /// things ride on it being IN the document's transaction rather than before
229 : /// it: the tenant row and the rows it accounts for commit together, and the
230 : /// claim holds a row lock the purge's `SELECT … FOR UPDATE` conflicts with,
231 : /// so the purge waits for an in-flight write instead of stepping over it.
232 : /// Committed separately, a write racing a purge leaves rows behind for a
233 : /// tenant the inventory no longer names — readable to whoever sends that
234 : /// tenant header, listed by nothing, and reclaimable by no further purge.
235 : ///
236 : /// The lock is SHARED. An upsert whose `DO UPDATE` re-set the tenant id took
237 : /// the exclusive lock and wrote a row version on every write, so writes in
238 : /// one tenant serialised behind each other and a table with one row per
239 : /// tenant collected one dead tuple per write in the broker. `FOR SHARE`
240 : /// conflicts with `FOR UPDATE` exactly as the exclusive lock did, and is
241 : /// compatible with itself, so concurrent writers hold it together.
242 : ///
243 : /// Idempotent, so two concurrent first writes both succeed (vs Scorpio's
244 : /// CREATE DATABASE + Flyway deadlock).
245 4427 : pub async fn claim_tenant(
246 4427 : tx: &mut Transaction<'_, Postgres>,
247 4427 : tenant: &TenantId,
248 4427 : ) -> Result<(), sqlx::Error> {
249 : // ponytail: shared row lock, so many concurrent writers in one tenant
250 : // register on one row; if multixact contention on that row ever shows up
251 : // in `pg_stat_activity`, the upgrade is a shared advisory lock keyed by
252 : // the tenant, which the purge takes exclusively and which touches no heap
253 : // page at all.
254 4427 : let claimed =
255 4427 : sqlx::query_scalar::<_, i32>("SELECT 1 FROM tenants WHERE tenant_id = $1 FOR SHARE")
256 4427 : .bind(tenant.as_str())
257 4427 : .fetch_optional(&mut **tx)
258 4427 : .await?
259 4427 : .is_some();
260 4427 : if claimed {
261 4389 : return Ok(());
262 38 : }
263 : // First write for this tenant. `DO UPDATE` rather than `DO NOTHING`:
264 : // a concurrent first write must block here until the winner commits,
265 : // and `DO NOTHING` would return without taking its lock.
266 38 : sqlx::query(
267 38 : "INSERT INTO tenants (tenant_id) VALUES ($1) \
268 38 : ON CONFLICT (tenant_id) DO UPDATE SET tenant_id = EXCLUDED.tenant_id",
269 38 : )
270 38 : .bind(tenant.as_str())
271 38 : .execute(&mut **tx)
272 38 : .await
273 38 : .map(|_| ())
274 4427 : }
275 :
276 : /// The same claim outside a document transaction, for the callers that have
277 : /// no document to pair it with (test seeds, the maintenance paths).
278 81 : pub async fn ensure_tenant(pool: &PgPool, tenant: &TenantId) -> Result<(), sqlx::Error> {
279 81 : sqlx::query("INSERT INTO tenants (tenant_id) VALUES ($1) ON CONFLICT DO NOTHING")
280 81 : .bind(tenant.as_str())
281 81 : .execute(pool)
282 81 : .await
283 81 : .map(|_| ())
284 81 : }
285 :
286 : #[cfg(test)]
287 : mod tests {
288 : use super::{bypasses, is_schema_mismatch, migrate_enabled};
289 :
290 : /// The RLS probe answers a security gate: a probe that errors must read as
291 : /// "unsafe", never as "the role does not bypass RLS".
292 : #[test]
293 2 : fn rls_probe_fails_closed() {
294 2 : assert!(!bypasses(Ok(false)), "a role without BYPASSRLS is safe");
295 2 : assert!(bypasses(Ok(true)), "superuser/BYPASSRLS bypasses");
296 2 : assert!(
297 2 : bypasses(Err(sqlx::Error::PoolTimedOut)),
298 : "an unanswerable probe must not pass the gate"
299 : );
300 2 : assert!(bypasses(Err(sqlx::Error::RowNotFound)));
301 2 : }
302 :
303 : /// A database from another release is a permanent failure; a database
304 : /// that is not up yet is not. The boot path retries the second and dies on
305 : /// the first, so the classification decides which cause an operator is
306 : /// shown — the squash of the pre-1.0 migration set left a 0.1.0 database
307 : /// reporting "not reachable after 30 s" for a schema it would never match.
308 : #[test]
309 2 : fn a_history_from_another_release_is_permanent_and_a_booting_database_is_not() {
310 : use sqlx::migrate::MigrateError;
311 12 : for e in [
312 2 : MigrateError::VersionMissing(5),
313 2 : MigrateError::VersionMismatch(1),
314 2 : MigrateError::VersionNotPresent(9),
315 2 : MigrateError::VersionTooOld(1, 4),
316 2 : MigrateError::VersionTooNew(9, 4),
317 2 : MigrateError::Dirty(3),
318 2 : ] {
319 12 : let shown = e.to_string();
320 12 : assert!(
321 12 : is_schema_mismatch(&sqlx::Error::Migrate(Box::new(e))),
322 : "{shown} can never be resolved by waiting"
323 : );
324 : }
325 2 : assert!(
326 2 : !is_schema_mismatch(&sqlx::Error::PoolTimedOut),
327 : "a database still booting is worth another attempt"
328 : );
329 2 : assert!(!is_schema_mismatch(&sqlx::Error::RowNotFound));
330 2 : assert!(
331 2 : !is_schema_mismatch(&sqlx::Error::Migrate(Box::new(MigrateError::Execute(
332 2 : sqlx::Error::PoolTimedOut
333 2 : )))),
334 : "a migration that failed to execute may be a lock the next attempt gets"
335 : );
336 2 : }
337 :
338 : /// Migrations stay on unless a deployment explicitly turns them off.
339 : #[test]
340 2 : fn migrate_switch_defaults_on() {
341 2 : assert!(migrate_enabled(None));
342 2 : assert!(migrate_enabled(Some("1")));
343 2 : assert!(migrate_enabled(Some("true")));
344 2 : assert!(!migrate_enabled(Some("0")));
345 2 : assert!(!migrate_enabled(Some("false")));
346 2 : }
347 : }
|
