agent_effects_postgres/
lib.rs

1//! PostgreSQL store for [`agent-effects`](https://crates.io/crates/agent-effects).
2//!
3//! For production services: many workers on many hosts sharing one database.
4//!
5//! ```no_run
6//! # async fn demo() -> Result<(), agent_effects_store::StoreError> {
7//! let store = agent_effects_postgres::PostgresStore::connect("postgres://localhost/effects").await?;
8//! # Ok(())
9//! # }
10//! ```
11//!
12//! # Sharing an application's database
13//!
14//! [`PostgresOptions`] fits the store into a database an application already
15//! uses, including through PgBouncer:
16//!
17//! - [`schema`](PostgresOptions::schema) keeps the store's tables and its
18//!   migration history in a schema of their own, named in every query. The
19//!   application's own `sqlx` migrations and `search_path` are left alone.
20//! - [`migrate`](PostgresOptions::migrate) and
21//!   [`attach`](PostgresOptions::attach) are separate steps: migrate once over
22//!   a direct connection, then attach to the application's pool without
23//!   migrating.
24//! - [`transaction_pooling`](PostgresOptions::transaction_pooling) makes every
25//!   query safe through a pooler in transaction mode.
26//!
27//! ```no_run
28//! # async fn demo(direct: sqlx::PgPool, pooled: sqlx::PgPool) -> Result<(), agent_effects_store::StoreError> {
29//! use agent_effects_postgres::PostgresOptions;
30//!
31//! let options = PostgresOptions::new()
32//!     .schema("agent_effects")
33//!     .transaction_pooling(true);
34//! // Once, at deploy or startup, over a direct connection to Postgres.
35//! options.migrate(&direct).await?;
36//! // The application's pool, which goes through PgBouncer.
37//! let store = options.attach(pooled).await?;
38//! # Ok(())
39//! # }
40//! ```
41//!
42//! # The database's clock
43//!
44//! By default ([`ClockSource::Database`]), every time the store uses comes
45//! from Postgres's `clock_timestamp()`, not from the calling worker:
46//!
47//! - lease liveness, lease expiry and takeover;
48//! - the times recorded on transitions;
49//! - the "no live lease" filter of recovery scans.
50//!
51//! Workers whose clocks disagree therefore still agree on who holds a lease.
52//! This removes the clock-skew limit that leases otherwise have. Retry
53//! schedules (`next_attempt_at`) are still computed by the worker.
54//! [`ClockSource::Caller`] uses the times the runtime passes in instead, as
55//! the in-memory and SQLite stores do; the conformance suite needs it to
56//! drive time.
57//!
58//! # Concurrency
59//!
60//! Every change is "load the row `FOR UPDATE`, apply the pure
61//! `EffectRecord` operation, save" in one transaction, so concurrent writers
62//! serialize on the row. The update also checks the version it read.
63//! Recovery and pending scans read with `FOR UPDATE SKIP LOCKED`: rows
64//! another worker is changing right now are skipped, not waited on.
65
66use std::ops::{Deref, DerefMut};
67use std::sync::Arc;
68use std::time::{Duration, SystemTime, UNIX_EPOCH};
69
70use agent_effects_store::{
71    EffectEvent, EffectId, EffectKey, EffectKind, EffectName, EffectRecord, EffectStatus,
72    EffectStore, ErrorRecord, InsertOutcome, Lease, ListQuery, LogicalKey, NewEffect, PruneQuery,
73    StoreError, Transition, TransitionRequest, WorkerId,
74};
75use serde_json::Value;
76use sqlx::migrate::Migrator;
77use sqlx::pool::PoolConnection;
78use sqlx::postgres::{PgArguments, PgConnectOptions, PgConnection, PgPool, PgPoolOptions, PgRow};
79use sqlx::query::Query;
80use sqlx::{AssertSqlSafe, Connection, Executor, Postgres, Row, Transaction};
81use time::OffsetDateTime;
82
83static MIGRATOR: Migrator = sqlx::migrate!("./migrations");
84
85/// sqlx's default name for the migration history table.
86const HISTORY: &str = "_sqlx_migrations";
87
88const COLUMNS: &str = "id, effect_name, logical_key, kind, status, input, input_fingerprint, \
89     output, last_error, created_by, attempt_count, may_have_applied, compensation_attempts, \
90     approved, next_attempt_at, attempt_started_at, attempt_ended_at, lease_owner, lease_epoch, \
91     lease_expires_at, version, created_at, updated_at, committed_at";
92
93/// Whose clock the store trusts for leases and recorded times.
94#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
95pub enum ClockSource {
96    /// Postgres's `clock_timestamp()`: one clock for every worker.
97    #[default]
98    Database,
99    /// The times the runtime passes in, from each worker's own clock.
100    Caller,
101}
102
103/// Where the store keeps its tables and how it talks to the database.
104///
105/// The defaults are what [`PostgresStore::connect`] uses: tables found
106/// through `search_path`, migration history in `_sqlx_migrations`, prepared
107/// statements cached per connection, and 10 connections.
108#[derive(Clone, Debug)]
109pub struct PostgresOptions {
110    schema: Option<String>,
111    transaction_pooling: bool,
112    max_connections: u32,
113}
114
115impl Default for PostgresOptions {
116    fn default() -> Self {
117        Self {
118            schema: None,
119            transaction_pooling: false,
120            max_connections: 10,
121        }
122    }
123}
124
125impl PostgresOptions {
126    /// The defaults.
127    pub fn new() -> Self {
128        Self::default()
129    }
130
131    /// Keeps the store's tables in `schema` and names it in every query, so
132    /// `search_path` does not matter. The migration history goes in
133    /// `schema._sqlx_migrations`, apart from any other `sqlx` migrations in
134    /// the database.
135    ///
136    /// `schema` is a name, used exactly as written (case included), not SQL.
137    /// [`migrate`](Self::migrate) creates it if it does not exist.
138    #[must_use]
139    pub fn schema(mut self, schema: impl Into<String>) -> Self {
140        self.schema = Some(schema.into());
141        self
142    }
143
144    /// For a pool that reaches Postgres through PgBouncer, or another pooler,
145    /// in transaction mode.
146    ///
147    /// Such a pooler may hand each transaction, and each statement outside
148    /// one, to a different server connection. So no statement is prepared
149    /// for reuse, and every store operation runs inside a transaction, so
150    /// that a statement's parse, bind and execution all reach the same
151    /// server. Reads that would otherwise run alone cost a `BEGIN` and a
152    /// `COMMIT` more.
153    #[must_use]
154    pub fn transaction_pooling(mut self, enabled: bool) -> Self {
155        self.transaction_pooling = enabled;
156        self
157    }
158
159    /// The size of the pool [`connect`](Self::connect) and
160    /// [`connect_with`](Self::connect_with) open; 10 by default.
161    /// [`attach`](Self::attach) uses the pool it is given as it is.
162    #[must_use]
163    pub fn max_connections(mut self, max: u32) -> Self {
164        self.max_connections = max;
165        self
166    }
167
168    /// Connects to `url`, applies migrations, and attaches.
169    ///
170    /// # Errors
171    ///
172    /// [`StoreError::Backend`] if connecting or migrating fails.
173    pub async fn connect(&self, url: &str) -> Result<PostgresStore, StoreError> {
174        let options: PgConnectOptions = url.parse().map_err(StoreError::backend)?;
175        self.connect_with(options).await
176    }
177
178    /// Connects with `options`, applies migrations, and attaches.
179    ///
180    /// # Errors
181    ///
182    /// [`StoreError::Backend`] if connecting or migrating fails.
183    pub async fn connect_with(
184        &self,
185        options: PgConnectOptions,
186    ) -> Result<PostgresStore, StoreError> {
187        let pool = PgPoolOptions::new()
188            .max_connections(self.max_connections)
189            .connect_with(options)
190            .await
191            .map_err(StoreError::backend)?;
192        self.migrate(&pool).await?;
193        self.attach(pool).await
194    }
195
196    /// Applies the store's migrations through `pool`.
197    ///
198    /// `pool` must reach Postgres directly, not through a pooler in
199    /// transaction mode: migrations hold a session-level advisory lock, so
200    /// concurrent callers are safe. With a [schema](Self::schema), the
201    /// connection used is closed afterwards instead of going back to the
202    /// pool, so `pool` keeps its own `search_path`.
203    ///
204    /// # Errors
205    ///
206    /// [`StoreError::Backend`] if migrating fails, or if the history table
207    /// holds a different migration under one of the store's versions (as
208    /// when it is shared with an application's migrations; give the store a
209    /// schema of its own).
210    pub async fn migrate(&self, pool: &PgPool) -> Result<(), StoreError> {
211        let Some(schema) = &self.schema else {
212            return MIGRATOR.run(pool).await.map_err(StoreError::backend);
213        };
214        let mut conn = pool.acquire().await.map_err(StoreError::backend)?.detach();
215        let result = migrate_into(&mut conn, schema).await;
216        // A failed close leaves only a connection the server will drop.
217        let _ = conn.close().await;
218        result
219    }
220
221    /// Uses `pool` without migrating; `pool` may go through a pooler.
222    ///
223    /// Checks that every migration this version needs has been applied, so a
224    /// missing [`migrate`](Self::migrate) fails here rather than on the first
225    /// effect.
226    ///
227    /// # Errors
228    ///
229    /// [`StoreError::Backend`] if the database is unreachable or a migration
230    /// is missing or not the store's own.
231    pub async fn attach(&self, pool: PgPool) -> Result<PostgresStore, StoreError> {
232        let store = PostgresStore {
233            pool,
234            clock: ClockSource::Database,
235            sql: Arc::new(Statements::new(self)),
236        };
237        store.check_migrated().await?;
238        Ok(store)
239    }
240}
241
242/// Runs the migrations on `conn` with `schema` as its only `search_path`, so
243/// the unqualified names in the migration files land in it.
244async fn migrate_into(conn: &mut PgConnection, schema: &str) -> Result<(), StoreError> {
245    let quoted = quote_ident(schema);
246    // CREATE SCHEMA IF NOT EXISTS needs CREATE on the database even when the
247    // schema exists, so leave a schema an administrator made alone.
248    let exists: bool =
249        sqlx::query_scalar("SELECT EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = $1)")
250            .bind(schema)
251            .fetch_one(&mut *conn)
252            .await
253            .map_err(StoreError::backend)?;
254    if !exists {
255        conn.execute(AssertSqlSafe(format!(
256            "CREATE SCHEMA IF NOT EXISTS {quoted}"
257        )))
258        .await
259        .map_err(StoreError::backend)?;
260    }
261    conn.execute(AssertSqlSafe(format!("SET search_path TO {quoted}")))
262        .await
263        .map_err(StoreError::backend)?;
264    let mut migrator = Migrator::with_migrations(MIGRATOR.iter().cloned().collect());
265    migrator.dangerous_set_table_name(format!("{quoted}.{HISTORY}"));
266    migrator.run(conn).await.map_err(StoreError::backend)
267}
268
269/// `name` as a quoted SQL identifier.
270fn quote_ident(name: &str) -> String {
271    format!("\"{}\"", name.replace('"', "\"\""))
272}
273
274/// The store's SQL, with its table names filled in.
275#[derive(Debug)]
276struct Statements {
277    transaction_pooling: bool,
278    history: String,
279    effects: String,
280    events: String,
281    lock: Arc<str>,
282    insert: Arc<str>,
283    get: Arc<str>,
284    get_by_key: Arc<str>,
285    save: Arc<str>,
286    insert_event: Arc<str>,
287    list_events: Arc<str>,
288}
289
290impl Statements {
291    fn new(options: &PostgresOptions) -> Self {
292        let qualify = |table: &str| match &options.schema {
293            Some(schema) => format!("{}.{table}", quote_ident(schema)),
294            None => table.to_owned(),
295        };
296        let (history, effects, events) = (
297            qualify(HISTORY),
298            qualify("effects"),
299            qualify("effect_events"),
300        );
301        Self {
302            transaction_pooling: options.transaction_pooling,
303            lock: format!("SELECT {COLUMNS} FROM {effects} WHERE id = $1 FOR UPDATE").into(),
304            insert: format!(
305                "INSERT INTO {effects} ({COLUMNS}) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, \
306                 $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24) \
307                 ON CONFLICT (effect_name, logical_key) DO NOTHING"
308            )
309            .into(),
310            get: format!("SELECT {COLUMNS} FROM {effects} WHERE id = $1").into(),
311            get_by_key: format!(
312                "SELECT {COLUMNS} FROM {effects} WHERE effect_name = $1 AND logical_key = $2"
313            )
314            .into(),
315            save: format!(
316                "UPDATE {effects} SET status = $1, output = $2, last_error = $3, \
317                 attempt_count = $4, may_have_applied = $5, compensation_attempts = $6, \
318                 approved = $7, next_attempt_at = $8, attempt_started_at = $9, \
319                 attempt_ended_at = $10, lease_owner = $11, lease_epoch = $12, \
320                 lease_expires_at = $13, version = $14, updated_at = $15, committed_at = $16 \
321                 WHERE id = $17 AND version = $18"
322            )
323            .into(),
324            insert_event: format!(
325                "INSERT INTO {events} \
326                 (effect_id, sequence, transition, from_status, to_status, attempt, actor, \
327                 payload, at) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)"
328            )
329            .into(),
330            list_events: format!(
331                "SELECT effect_id, sequence, transition, from_status, to_status, attempt, \
332                 actor, payload, at FROM {events} WHERE effect_id = $1 ORDER BY sequence"
333            )
334            .into(),
335            history,
336            effects,
337            events,
338        }
339    }
340
341    /// A query for `sql`, prepared for reuse unless behind a pooler.
342    fn query<'q>(&self, sql: impl Into<Arc<str>>) -> Query<'q, Postgres, PgArguments> {
343        sqlx::query(AssertSqlSafe(sql.into())).persistent(!self.transaction_pooling)
344    }
345}
346
347/// The connection one store operation runs on: in a transaction when the
348/// pool goes through a transaction-mode pooler, so every statement of the
349/// operation reaches the same server.
350enum Session {
351    Direct(PoolConnection<Postgres>),
352    Pinned(Transaction<'static, Postgres>),
353}
354
355impl Session {
356    /// Ends the operation, committing its transaction if it has one.
357    async fn finish(self) -> Result<(), StoreError> {
358        match self {
359            Self::Direct(_) => Ok(()),
360            Self::Pinned(tx) => tx.commit().await.map_err(StoreError::backend),
361        }
362    }
363}
364
365impl Deref for Session {
366    type Target = PgConnection;
367
368    fn deref(&self) -> &PgConnection {
369        match self {
370            Self::Direct(conn) => conn,
371            Self::Pinned(tx) => tx,
372        }
373    }
374}
375
376impl DerefMut for Session {
377    fn deref_mut(&mut self) -> &mut PgConnection {
378        match self {
379            Self::Direct(conn) => conn,
380            Self::Pinned(tx) => tx,
381        }
382    }
383}
384
385/// An [`EffectStore`] in a PostgreSQL database.
386///
387/// Cheap to clone; clones share the connection pool.
388#[derive(Clone, Debug)]
389pub struct PostgresStore {
390    pool: PgPool,
391    clock: ClockSource,
392    sql: Arc<Statements>,
393}
394
395impl PostgresStore {
396    /// Connects to `url` and applies migrations, with the default
397    /// [`PostgresOptions`].
398    ///
399    /// # Errors
400    ///
401    /// [`StoreError::Backend`] if connecting or migrating fails.
402    pub async fn connect(url: &str) -> Result<Self, StoreError> {
403        PostgresOptions::new().connect(url).await
404    }
405
406    /// Connects with `options` and applies migrations, with the default
407    /// [`PostgresOptions`]. For a schema of the store's own, see
408    /// [`PostgresOptions::schema`].
409    ///
410    /// # Errors
411    ///
412    /// [`StoreError::Backend`] if connecting or migrating fails.
413    pub async fn connect_with(options: PgConnectOptions) -> Result<Self, StoreError> {
414        PostgresOptions::new().connect_with(options).await
415    }
416
417    /// Uses an existing pool and applies migrations, with the default
418    /// [`PostgresOptions`]. Concurrent callers are safe: migrations run under
419    /// an advisory lock. To use a pool without migrating, as through
420    /// PgBouncer, see [`PostgresOptions::attach`].
421    ///
422    /// # Errors
423    ///
424    /// [`StoreError::Backend`] if migrating fails.
425    pub async fn from_pool(pool: PgPool) -> Result<Self, StoreError> {
426        let options = PostgresOptions::new();
427        options.migrate(&pool).await?;
428        options.attach(pool).await
429    }
430
431    /// Chooses whose clock to trust; see [`ClockSource`].
432    #[must_use]
433    pub fn with_clock_source(mut self, clock: ClockSource) -> Self {
434        self.clock = clock;
435        self
436    }
437
438    /// The underlying pool.
439    pub fn pool(&self) -> &PgPool {
440        &self.pool
441    }
442
443    async fn session(&self) -> Result<Session, StoreError> {
444        if self.sql.transaction_pooling {
445            self.pool.begin().await.map(Session::Pinned)
446        } else {
447            self.pool.acquire().await.map(Session::Direct)
448        }
449        .map_err(StoreError::backend)
450    }
451
452    /// Fails unless every migration is applied, and is the store's own.
453    async fn check_migrated(&self) -> Result<(), StoreError> {
454        let mut session = self.session().await?;
455        let applied = self
456            .sql
457            .query(format!(
458                "SELECT version, checksum FROM {} WHERE success",
459                self.sql.history
460            ))
461            .fetch_all(&mut *session)
462            .await;
463        let applied = match applied {
464            Ok(rows) => rows,
465            Err(e)
466                if e.as_database_error()
467                    .and_then(sqlx::error::DatabaseError::code)
468                    .as_deref()
469                    == Some("42P01") =>
470            {
471                return Err(StoreError::backend(format!(
472                    "no migration history at {}: run PostgresOptions::migrate first",
473                    self.sql.history
474                )));
475            }
476            Err(e) => return Err(StoreError::backend(e)),
477        };
478        session.finish().await?;
479        for migration in MIGRATOR.iter() {
480            if migration.migration_type.is_down_migration() {
481                continue;
482            }
483            let found = applied.iter().find_map(|row| {
484                (row.try_get::<i64, _>("version").ok()? == migration.version)
485                    .then(|| row.try_get::<Vec<u8>, _>("checksum").ok())
486            });
487            match found {
488                None => {
489                    return Err(StoreError::backend(format!(
490                        "migration {} ({}) is not applied in {}: run PostgresOptions::migrate \
491                         first",
492                        migration.version, migration.description, self.sql.history
493                    )));
494                }
495                Some(checksum) if checksum.as_deref() != Some(&*migration.checksum) => {
496                    return Err(StoreError::backend(format!(
497                        "migration {} in {} is not agent-effects' own; give the store a \
498                         schema of its own with PostgresOptions::schema",
499                        migration.version, self.sql.history
500                    )));
501                }
502                Some(_) => {}
503            }
504        }
505        Ok(())
506    }
507
508    async fn database_now(&self, conn: &mut PgConnection) -> Result<SystemTime, StoreError> {
509        let now: OffsetDateTime = sqlx::query_scalar("SELECT clock_timestamp()")
510            .persistent(!self.sql.transaction_pooling)
511            .fetch_one(conn)
512            .await
513            .map_err(StoreError::backend)?;
514        Ok(SystemTime::from(now))
515    }
516
517    /// The time to use: the database's when it is the source, else `caller`.
518    fn effective(&self, database: SystemTime, caller: SystemTime) -> SystemTime {
519        match self.clock {
520            ClockSource::Database => database,
521            ClockSource::Caller => caller,
522        }
523    }
524
525    async fn fetch_by_key(
526        &self,
527        conn: &mut PgConnection,
528        key: &EffectKey,
529    ) -> Result<Option<EffectRecord>, StoreError> {
530        self.sql
531            .query(Arc::clone(&self.sql.get_by_key))
532            .bind(key.name.as_str())
533            .bind(key.key.as_str())
534            .fetch_optional(conn)
535            .await
536            .map_err(StoreError::backend)?
537            .as_ref()
538            .map(decode_record)
539            .transpose()
540    }
541
542    async fn fetch(
543        &self,
544        conn: &mut PgConnection,
545        id: EffectId,
546    ) -> Result<Option<EffectRecord>, StoreError> {
547        self.sql
548            .query(Arc::clone(&self.sql.get))
549            .bind(*id.as_uuid())
550            .fetch_optional(conn)
551            .await
552            .map_err(StoreError::backend)?
553            .as_ref()
554            .map(decode_record)
555            .transpose()
556    }
557
558    /// Loads `id` locked `FOR UPDATE`, applies `change` (given the
559    /// database's current time once the lock is held), and saves the
560    /// record. The transaction stays open for further writes.
561    async fn modify<T>(
562        &self,
563        id: EffectId,
564        change: impl FnOnce(&mut EffectRecord, SystemTime) -> Result<T, StoreError>,
565    ) -> Result<(T, EffectRecord, Transaction<'static, Postgres>), StoreError> {
566        let mut tx = self.pool.begin().await.map_err(StoreError::backend)?;
567        let row = self
568            .sql
569            .query(Arc::clone(&self.sql.lock))
570            .bind(*id.as_uuid())
571            .fetch_optional(&mut *tx)
572            .await
573            .map_err(StoreError::backend)?
574            .ok_or(StoreError::NotFound(id))?;
575        let mut record = decode_record(&row)?;
576        // Read only now: a clock read in the locking statement is taken
577        // before waiting for the lock, and can be stale by the whole wait.
578        let db_now = self.database_now(&mut tx).await?;
579        let read_version = record.version;
580        let result = change(&mut record, db_now)?;
581        normalize(&mut record);
582        self.save(&mut tx, &record, read_version).await?;
583        Ok((result, record, tx))
584    }
585
586    /// Writes `record` over the row last read at `read_version`.
587    async fn save(
588        &self,
589        tx: &mut Transaction<'static, Postgres>,
590        record: &EffectRecord,
591        read_version: u64,
592    ) -> Result<(), StoreError> {
593        let updated = self
594            .sql
595            .query(Arc::clone(&self.sql.save))
596            .bind(record.status.as_str())
597            .bind(record.output.clone())
598            .bind(json(record.last_error.as_ref())?)
599            .bind(to_i32(record.attempt_count)?)
600            .bind(record.may_have_applied)
601            .bind(to_i32(record.compensation_attempts)?)
602            .bind(record.approved)
603            .bind(record.next_attempt_at.map(timestamp))
604            .bind(record.attempt_started_at.map(timestamp))
605            .bind(record.attempt_ended_at.map(timestamp))
606            .bind(record.lease_owner.as_ref().map(|w| w.as_str().to_owned()))
607            .bind(to_i64(record.lease_epoch)?)
608            .bind(record.lease_expires_at.map(timestamp))
609            .bind(to_i64(record.version)?)
610            .bind(timestamp(record.updated_at))
611            .bind(record.committed_at.map(timestamp))
612            .bind(*record.id.as_uuid())
613            .bind(to_i64(read_version)?)
614            .execute(&mut **tx)
615            .await
616            .map_err(StoreError::backend)?
617            .rows_affected();
618        if updated == 1 {
619            Ok(())
620        } else {
621            // Unreachable under FOR UPDATE; kept as a second guard.
622            Err(StoreError::backend(format!(
623                "effect {} changed under a row lock",
624                record.id
625            )))
626        }
627    }
628}
629
630impl EffectStore for PostgresStore {
631    async fn insert_or_get(&self, mut new: NewEffect) -> Result<InsertOutcome, StoreError> {
632        let mut session = self.session().await?;
633        if self.clock == ClockSource::Database {
634            new.now = self.database_now(&mut session).await?;
635        }
636        let mut record = EffectRecord::new(new);
637        normalize(&mut record);
638        let inserted = self
639            .sql
640            .query(Arc::clone(&self.sql.insert))
641            .bind(*record.id.as_uuid())
642            .bind(record.key.name.as_str())
643            .bind(record.key.key.as_str())
644            .bind(record.kind.as_str())
645            .bind(record.status.as_str())
646            .bind(record.input.clone())
647            .bind(record.input_fingerprint.clone())
648            .bind(record.output.clone())
649            .bind(json(record.last_error.as_ref())?)
650            .bind(record.created_by.clone())
651            .bind(to_i32(record.attempt_count)?)
652            .bind(record.may_have_applied)
653            .bind(to_i32(record.compensation_attempts)?)
654            .bind(record.approved)
655            .bind(record.next_attempt_at.map(timestamp))
656            .bind(record.attempt_started_at.map(timestamp))
657            .bind(record.attempt_ended_at.map(timestamp))
658            .bind(record.lease_owner.as_ref().map(|w| w.as_str().to_owned()))
659            .bind(to_i64(record.lease_epoch)?)
660            .bind(record.lease_expires_at.map(timestamp))
661            .bind(to_i64(record.version)?)
662            .bind(timestamp(record.created_at))
663            .bind(timestamp(record.updated_at))
664            .bind(record.committed_at.map(timestamp))
665            .execute(&mut *session)
666            .await
667            .map_err(StoreError::backend)?
668            .rows_affected()
669            == 1;
670        if inserted {
671            session.finish().await?;
672            return Ok(InsertOutcome {
673                record,
674                inserted: true,
675            });
676        }
677        let existing = self
678            .fetch_by_key(&mut session, &record.key)
679            .await?
680            .ok_or_else(|| StoreError::backend("conflicting insert, but no existing record"))?;
681        session.finish().await?;
682        Ok(InsertOutcome {
683            record: existing,
684            inserted: false,
685        })
686    }
687
688    async fn get(&self, id: EffectId) -> Result<Option<EffectRecord>, StoreError> {
689        let mut session = self.session().await?;
690        let record = self.fetch(&mut session, id).await?;
691        session.finish().await?;
692        Ok(record)
693    }
694
695    async fn get_by_key(&self, key: &EffectKey) -> Result<Option<EffectRecord>, StoreError> {
696        let mut session = self.session().await?;
697        let record = self.fetch_by_key(&mut session, key).await?;
698        session.finish().await?;
699        Ok(record)
700    }
701
702    async fn acquire_lease(
703        &self,
704        id: EffectId,
705        owner: &WorkerId,
706        now: SystemTime,
707        ttl: Duration,
708    ) -> Result<Lease, StoreError> {
709        let (mut lease, _, tx) = self
710            .modify(id, |record, db_now| {
711                record.acquire_lease(owner, self.effective(db_now, now), ttl)
712            })
713            .await?;
714        tx.commit().await.map_err(StoreError::backend)?;
715        lease.expires_at = round_trip(lease.expires_at);
716        Ok(lease)
717    }
718
719    async fn renew_lease(
720        &self,
721        lease: &Lease,
722        now: SystemTime,
723        ttl: Duration,
724    ) -> Result<Lease, StoreError> {
725        let (mut renewed, _, tx) = self
726            .modify(lease.effect_id, |record, db_now| {
727                record.renew_lease(lease, self.effective(db_now, now), ttl)
728            })
729            .await?;
730        tx.commit().await.map_err(StoreError::backend)?;
731        renewed.expires_at = round_trip(renewed.expires_at);
732        Ok(renewed)
733    }
734
735    async fn release_lease(&self, lease: &Lease) -> Result<(), StoreError> {
736        let (released, _, tx) = self
737            .modify(lease.effect_id, |record, _| Ok(record.release_lease(lease)))
738            .await?;
739        if released {
740            tx.commit().await.map_err(StoreError::backend)?;
741        } else {
742            tx.rollback().await.map_err(StoreError::backend)?;
743        }
744        Ok(())
745    }
746
747    async fn transition(&self, mut request: TransitionRequest) -> Result<EffectRecord, StoreError> {
748        let id = request.id;
749        let (mut event, record, mut tx) = self
750            .modify(id, |record, db_now| {
751                request.now = self.effective(db_now, request.now);
752                record.apply(request)
753            })
754            .await?;
755        event.at = round_trip(event.at);
756        self.sql
757            .query(Arc::clone(&self.sql.insert_event))
758            .bind(*event.effect_id.as_uuid())
759            .bind(to_i64(event.sequence)?)
760            .bind(event.transition.as_str())
761            .bind(event.from.as_str())
762            .bind(event.to.as_str())
763            .bind(to_i32(event.attempt)?)
764            .bind(event.actor.clone())
765            .bind(event.payload.clone())
766            .bind(timestamp(event.at))
767            .execute(&mut *tx)
768            .await
769            .map_err(StoreError::backend)?;
770        tx.commit().await.map_err(StoreError::backend)?;
771        Ok(record)
772    }
773
774    async fn list(&self, query: ListQuery) -> Result<Vec<EffectRecord>, StoreError> {
775        let mut sql = format!("SELECT {COLUMNS} FROM {} WHERE TRUE", self.sql.effects);
776        let mut next = 1;
777        let mut placeholder = || {
778            let p = format!("${next}");
779            next += 1;
780            p
781        };
782        if !query.statuses.is_empty() {
783            let marks: Vec<String> = query.statuses.iter().map(|_| placeholder()).collect();
784            sql.push_str(" AND status IN (");
785            sql.push_str(&marks.join(", "));
786            sql.push(')');
787        }
788        let lease_filter = query.lease_expired_at.is_some();
789        if lease_filter {
790            let now = match self.clock {
791                ClockSource::Database => "clock_timestamp()".to_owned(),
792                ClockSource::Caller => placeholder(),
793            };
794            sql.push_str(
795                " AND (lease_owner IS NULL OR lease_expires_at IS NULL OR lease_expires_at <= ",
796            );
797            sql.push_str(&now);
798            sql.push(')');
799        }
800        if query.after.is_some() {
801            sql.push_str(" AND id > ");
802            sql.push_str(&placeholder());
803        }
804        sql.push_str(" ORDER BY id LIMIT ");
805        sql.push_str(&placeholder());
806        if lease_filter {
807            // A work-finding scan: skip rows another worker is changing.
808            sql.push_str(" FOR UPDATE SKIP LOCKED");
809        }
810
811        // Only placeholders and the table name are interpolated; every value
812        // is bound.
813        let mut statement = self.sql.query(sql);
814        for status in &query.statuses {
815            statement = statement.bind(status.as_str());
816        }
817        if let (Some(now), ClockSource::Caller) = (query.lease_expired_at, self.clock) {
818            statement = statement.bind(timestamp(now));
819        }
820        if let Some(after) = query.after {
821            statement = statement.bind(*after.as_uuid());
822        }
823        statement = statement.bind(i64::try_from(query.limit).unwrap_or(i64::MAX));
824        let mut session = self.session().await?;
825        let records = statement
826            .fetch_all(&mut *session)
827            .await
828            .map_err(StoreError::backend)?
829            .iter()
830            .map(decode_record)
831            .collect::<Result<_, _>>()?;
832        session.finish().await?;
833        Ok(records)
834    }
835
836    async fn events(&self, id: EffectId) -> Result<Vec<EffectEvent>, StoreError> {
837        let mut session = self.session().await?;
838        if self.fetch(&mut session, id).await?.is_none() {
839            return Err(StoreError::NotFound(id));
840        }
841        let events = self
842            .sql
843            .query(Arc::clone(&self.sql.list_events))
844            .bind(*id.as_uuid())
845            .fetch_all(&mut *session)
846            .await
847            .map_err(StoreError::backend)?
848            .iter()
849            .map(decode_event)
850            .collect::<Result<_, _>>()?;
851        session.finish().await?;
852        Ok(events)
853    }
854
855    async fn prune(&self, query: PruneQuery) -> Result<u64, StoreError> {
856        let Some(cutoff) = query.cutoff().filter(|_| query.status.is_settled()) else {
857            return Ok(0);
858        };
859        // Rows another worker is changing are skipped, and the ones selected
860        // stay locked until both deletes commit.
861        let (cutoff_sql, now_sql, limit) = match self.clock {
862            ClockSource::Database => (
863                "clock_timestamp() - $2 * interval '1 microsecond'",
864                "clock_timestamp()",
865                "$3",
866            ),
867            ClockSource::Caller => ("$2", "$3", "$4"),
868        };
869        let (effects, events) = (&self.sql.effects, &self.sql.events);
870        let sql = format!(
871            "WITH doomed AS (\
872                SELECT id FROM {effects} WHERE status = $1 AND updated_at <= {cutoff_sql} \
873                AND (lease_owner IS NULL OR lease_expires_at IS NULL \
874                     OR lease_expires_at <= {now_sql}) \
875                ORDER BY id LIMIT {limit} FOR UPDATE SKIP LOCKED), \
876             events AS (DELETE FROM {events} WHERE effect_id IN (SELECT id FROM doomed)) \
877             DELETE FROM {effects} WHERE id IN (SELECT id FROM doomed)"
878        );
879        // Only fixed fragments and table names are interpolated; every value
880        // is bound.
881        let mut statement = self.sql.query(sql).bind(query.status.as_str());
882        statement = match self.clock {
883            ClockSource::Database => statement
884                .bind(i64::try_from(query.older_than.as_micros()).map_err(StoreError::backend)?),
885            ClockSource::Caller => statement.bind(timestamp(cutoff)).bind(timestamp(query.now)),
886        };
887        let mut session = self.session().await?;
888        let deleted = statement
889            .bind(i64::try_from(query.limit).unwrap_or(i64::MAX))
890            .execute(&mut *session)
891            .await
892            .map_err(StoreError::backend)?;
893        session.finish().await?;
894        Ok(deleted.rows_affected())
895    }
896}
897
898fn decode_record(row: &PgRow) -> Result<EffectRecord, StoreError> {
899    let kind: String = get(row, "kind")?;
900    let status: String = get(row, "status")?;
901    Ok(EffectRecord {
902        id: EffectId::from_uuid(get(row, "id")?),
903        key: EffectKey::new(
904            EffectName::new(get::<String>(row, "effect_name")?).map_err(StoreError::backend)?,
905            LogicalKey::new(get::<String>(row, "logical_key")?).map_err(StoreError::backend)?,
906        ),
907        kind: EffectKind::parse(&kind).ok_or_else(|| corrupt("kind", &kind))?,
908        status: EffectStatus::parse(&status).ok_or_else(|| corrupt("status", &status))?,
909        input: get(row, "input")?,
910        input_fingerprint: get(row, "input_fingerprint")?,
911        output: get(row, "output")?,
912        last_error: get::<Option<Value>>(row, "last_error")?
913            .map(serde_json::from_value::<ErrorRecord>)
914            .transpose()
915            .map_err(StoreError::backend)?,
916        created_by: get(row, "created_by")?,
917        attempt_count: to_u32(get(row, "attempt_count")?)?,
918        may_have_applied: get(row, "may_have_applied")?,
919        compensation_attempts: to_u32(get(row, "compensation_attempts")?)?,
920        approved: get(row, "approved")?,
921        next_attempt_at: get::<Option<OffsetDateTime>>(row, "next_attempt_at")?
922            .map(SystemTime::from),
923        attempt_started_at: get::<Option<OffsetDateTime>>(row, "attempt_started_at")?
924            .map(SystemTime::from),
925        attempt_ended_at: get::<Option<OffsetDateTime>>(row, "attempt_ended_at")?
926            .map(SystemTime::from),
927        lease_owner: get::<Option<String>>(row, "lease_owner")?.map(WorkerId::new),
928        lease_epoch: to_u64(get(row, "lease_epoch")?)?,
929        lease_expires_at: get::<Option<OffsetDateTime>>(row, "lease_expires_at")?
930            .map(SystemTime::from),
931        version: to_u64(get(row, "version")?)?,
932        created_at: SystemTime::from(get::<OffsetDateTime>(row, "created_at")?),
933        updated_at: SystemTime::from(get::<OffsetDateTime>(row, "updated_at")?),
934        committed_at: get::<Option<OffsetDateTime>>(row, "committed_at")?.map(SystemTime::from),
935    })
936}
937
938fn decode_event(row: &PgRow) -> Result<EffectEvent, StoreError> {
939    let transition: String = get(row, "transition")?;
940    let from: String = get(row, "from_status")?;
941    let to: String = get(row, "to_status")?;
942    Ok(EffectEvent {
943        effect_id: EffectId::from_uuid(get(row, "effect_id")?),
944        sequence: to_u64(get(row, "sequence")?)?,
945        transition: Transition::parse(&transition)
946            .ok_or_else(|| corrupt("transition", &transition))?,
947        from: EffectStatus::parse(&from).ok_or_else(|| corrupt("from_status", &from))?,
948        to: EffectStatus::parse(&to).ok_or_else(|| corrupt("to_status", &to))?,
949        attempt: to_u32(get(row, "attempt")?)?,
950        actor: get(row, "actor")?,
951        payload: get(row, "payload")?,
952        at: SystemTime::from(get::<OffsetDateTime>(row, "at")?),
953    })
954}
955
956fn get<'r, T>(row: &'r PgRow, column: &str) -> Result<T, StoreError>
957where
958    T: sqlx::Decode<'r, Postgres> + sqlx::Type<Postgres>,
959{
960    row.try_get(column).map_err(StoreError::backend)
961}
962
963fn corrupt(column: &str, value: &str) -> StoreError {
964    StoreError::backend(format!(
965        "unreadable {column} in effects database: {value:?}"
966    ))
967}
968
969fn json(error: Option<&ErrorRecord>) -> Result<Option<Value>, StoreError> {
970    error
971        .map(serde_json::to_value)
972        .transpose()
973        .map_err(StoreError::backend)
974}
975
976fn to_i32(value: u32) -> Result<i32, StoreError> {
977    i32::try_from(value).map_err(StoreError::backend)
978}
979
980fn to_u32(value: i32) -> Result<u32, StoreError> {
981    u32::try_from(value).map_err(StoreError::backend)
982}
983
984fn to_i64(value: u64) -> Result<i64, StoreError> {
985    i64::try_from(value).map_err(StoreError::backend)
986}
987
988fn to_u64(value: i64) -> Result<u64, StoreError> {
989    u64::try_from(value).map_err(StoreError::backend)
990}
991
992fn timestamp(time: SystemTime) -> OffsetDateTime {
993    OffsetDateTime::from(round_trip(time))
994}
995
996/// `time` at millisecond precision, the precision the store keeps.
997fn round_trip(time: SystemTime) -> SystemTime {
998    match time.duration_since(UNIX_EPOCH) {
999        Ok(after) => {
1000            UNIX_EPOCH + Duration::from_millis(u64::try_from(after.as_millis()).unwrap_or(u64::MAX))
1001        }
1002        Err(before) => {
1003            let millis = u64::try_from(before.duration().as_millis()).unwrap_or(u64::MAX);
1004            UNIX_EPOCH - Duration::from_millis(millis)
1005        }
1006    }
1007}
1008
1009/// Gives every time in `record` its stored precision, so returned records
1010/// equal what a later read returns.
1011fn normalize(record: &mut EffectRecord) {
1012    for time in [
1013        &mut record.next_attempt_at,
1014        &mut record.attempt_started_at,
1015        &mut record.attempt_ended_at,
1016        &mut record.lease_expires_at,
1017        &mut record.committed_at,
1018    ]
1019    .into_iter()
1020    .flatten()
1021    {
1022        *time = round_trip(*time);
1023    }
1024    record.created_at = round_trip(record.created_at);
1025    record.updated_at = round_trip(record.updated_at);
1026}
agent_effects_postgres/lib.rs

agent_effects_postgres/
lib.rs
