Skip to main content

cratestack_sqlx/
idempotency.rs

1//! Postgres-backed [`IdempotencyStore`]. Banks need duplicate-execution
2//! protection even under concurrency, so this uses the atomic-reservation
3//! pattern: a single upsert claims the key (or surfaces the existing
4//! claim), the middleware runs the handler only when it owns the
5//! reservation, then writes the captured response with `complete`.
6//!
7//! Expired rows are replaced on the next reservation via the
8//! `ON CONFLICT DO UPDATE WHERE cratestack_idempotency.expires_at <= NOW()`
9//! clause, letting the new caller take over a stale row in the same
10//! statement that would otherwise have hit the unique-key wall.
11
12mod operations;
13
14use std::time::SystemTime;
15
16use async_trait::async_trait;
17use cratestack_axum::idempotency::{IDEMPOTENCY_TABLE_DDL, IdempotencyStore, ReservationOutcome};
18use cratestack_core::CoolError;
19
20use crate::sqlx;
21
22#[derive(Clone)]
23pub struct SqlxIdempotencyStore {
24    pool: sqlx::PgPool,
25}
26
27impl SqlxIdempotencyStore {
28    pub fn new(pool: sqlx::PgPool) -> Self {
29        Self { pool }
30    }
31
32    /// Ensure the table exists. Banks typically run this via their own
33    /// migration tooling; exposed here for convenience.
34    pub async fn ensure_schema(&self) -> Result<(), CoolError> {
35        // Multi-statement DDL (table + index) — `raw_sql` sends it as one
36        // batch over PG's simple-query protocol instead of splitting on
37        // `;` client-side, which would corrupt any dollar-quoted body.
38        sqlx::raw_sql(IDEMPOTENCY_TABLE_DDL)
39            .execute(&self.pool)
40            .await
41            .map_err(|error| CoolError::Database(error.to_string()))?;
42        Ok(())
43    }
44
45    /// Delete expired rows. Run periodically — the request path does
46    /// not auto-GC, although `reserve_or_fetch` does take over any
47    /// single expired row it tries to claim.
48    pub async fn garbage_collect(&self) -> Result<u64, CoolError> {
49        let result = sqlx::query("DELETE FROM cratestack_idempotency WHERE expires_at < NOW()")
50            .execute(&self.pool)
51            .await
52            .map_err(|error| CoolError::Database(error.to_string()))?;
53        Ok(result.rows_affected())
54    }
55}
56
57#[async_trait]
58impl IdempotencyStore for SqlxIdempotencyStore {
59    async fn reserve_or_fetch(
60        &self,
61        principal: &str,
62        key: &str,
63        request_hash: [u8; 32],
64        expires_at: SystemTime,
65    ) -> Result<ReservationOutcome, CoolError> {
66        operations::reserve_or_fetch(&self.pool, principal, key, request_hash, expires_at).await
67    }
68
69    async fn complete(
70        &self,
71        principal: &str,
72        key: &str,
73        token: uuid::Uuid,
74        status: u16,
75        headers: &[u8],
76        body: &[u8],
77    ) -> Result<(), CoolError> {
78        operations::complete(&self.pool, principal, key, token, status, headers, body).await
79    }
80
81    async fn release(
82        &self,
83        principal: &str,
84        key: &str,
85        token: uuid::Uuid,
86    ) -> Result<(), CoolError> {
87        operations::release(&self.pool, principal, key, token).await
88    }
89}
90
91/// Compute when a record originally captured at `created_at` will
92/// expire. Pulled out for unit-test reach; the SystemTime arithmetic
93/// is otherwise awkward to assert against without a clock injection
94/// point.
95pub fn expiry_from(created_at: SystemTime, ttl: std::time::Duration) -> SystemTime {
96    created_at + ttl
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use std::time::{Duration, SystemTime};
103
104    #[test]
105    fn expiry_adds_ttl_to_creation() {
106        let now = SystemTime::UNIX_EPOCH;
107        let expiry = expiry_from(now, Duration::from_secs(24 * 3600));
108        assert_eq!(
109            expiry.duration_since(SystemTime::UNIX_EPOCH).unwrap(),
110            Duration::from_secs(24 * 3600),
111        );
112    }
113}