eventuary_postgres/
claim_buffer.rs

1//! Postgres-backed `ClaimedBufferStore` implementation.
2//!
3//! Atomic claim semantics via `FOR UPDATE SKIP LOCKED`: concurrent claim
4//! batches receive disjoint sets without explicit synchronization.
5
6use std::sync::Arc;
7use std::time::Duration;
8
9use chrono::Utc;
10use sqlx::{PgPool, Row};
11
12use eventuary_core::io::OwnerId;
13use eventuary_core::io::reader::claim_buffer::{ClaimedBufferEntry, ClaimedBufferStore};
14use eventuary_core::{Error, Event, Result, SerializedEvent};
15
16use crate::relation::PgRelationName;
17use crate::schema::{Migration, RelationReplacement};
18
19const CLAIMED_BUFFER_STORE_0001_INIT_SQL: &str = r#"
20CREATE TABLE IF NOT EXISTS {buffer_claims} (
21    id              BIGSERIAL    PRIMARY KEY,
22    event           JSONB        NOT NULL,
23    claimed_by      TEXT         NULL,
24    claimed_until   TIMESTAMPTZ  NULL,
25    attempts        INT          NOT NULL DEFAULT 0,
26    created_at      TIMESTAMPTZ  NOT NULL DEFAULT NOW()
27);
28
29CREATE INDEX IF NOT EXISTS idx_buffer_claims_pending
30ON {buffer_claims} (claimed_until, id)
31WHERE claimed_by IS NULL OR claimed_until IS NULL;
32
33CREATE INDEX IF NOT EXISTS idx_buffer_claims_visibility
34ON {buffer_claims} (claimed_until)
35WHERE claimed_by IS NOT NULL;
36"#;
37
38const CLAIMED_BUFFER_STORE_MIGRATIONS: &[Migration] = &[Migration {
39    name: "0001_init",
40    sql: CLAIMED_BUFFER_STORE_0001_INIT_SQL,
41}];
42
43#[derive(Debug, Clone)]
44pub struct PgClaimedBufferStoreConfig {
45    pub relation: PgRelationName,
46}
47
48impl Default for PgClaimedBufferStoreConfig {
49    fn default() -> Self {
50        Self {
51            relation: PgRelationName::new("event_buffer_claims")
52                .expect("default buffer claims relation"),
53        }
54    }
55}
56
57pub struct PgClaimedBufferStore {
58    pool: PgPool,
59    relation: Arc<String>,
60}
61
62impl Clone for PgClaimedBufferStore {
63    fn clone(&self) -> Self {
64        Self {
65            pool: self.pool.clone(),
66            relation: Arc::clone(&self.relation),
67        }
68    }
69}
70
71impl PgClaimedBufferStore {
72    pub fn new(pool: PgPool, config: PgClaimedBufferStoreConfig) -> Self {
73        Self {
74            pool,
75            relation: Arc::new(config.relation.render()),
76        }
77    }
78
79    pub async fn connect(pool: PgPool, config: PgClaimedBufferStoreConfig) -> Result<Self> {
80        Self::prepare_schema(&pool, &config).await?;
81        Ok(Self::new(pool, config))
82    }
83
84    pub async fn prepare_schema(pool: &PgPool, config: &PgClaimedBufferStoreConfig) -> Result<()> {
85        crate::schema::apply_schema(
86            pool,
87            CLAIMED_BUFFER_STORE_MIGRATIONS,
88            &[RelationReplacement {
89                token: "{buffer_claims}",
90                relation: &config.relation,
91            }],
92        )
93        .await
94    }
95
96    pub fn schema_sql(config: &PgClaimedBufferStoreConfig) -> String {
97        crate::schema::render_schema_sql(
98            CLAIMED_BUFFER_STORE_MIGRATIONS,
99            &[RelationReplacement {
100                token: "{buffer_claims}",
101                relation: &config.relation,
102            }],
103        )
104    }
105}
106
107fn encode_event(event: &Event) -> Result<serde_json::Value> {
108    let serialized = SerializedEvent::from_event(event)?;
109    serde_json::to_value(serialized)
110        .map_err(|e| Error::Serialization(format!("claim buffer encode event: {e}")))
111}
112
113fn decode_event(value: serde_json::Value) -> Result<Event> {
114    let serialized: SerializedEvent = serde_json::from_value(value)
115        .map_err(|e| Error::Serialization(format!("claim buffer decode event: {e}")))?;
116    serialized.to_event()
117}
118
119impl ClaimedBufferStore for PgClaimedBufferStore {
120    type Id = i64;
121
122    async fn push(&self, event: &Event) -> Result<Self::Id> {
123        let sql = format!(
124            "INSERT INTO {relation} (event) VALUES ($1) RETURNING id",
125            relation = self.relation
126        );
127        let row = sqlx::query(&sql)
128            .bind(encode_event(event)?)
129            .fetch_one(&self.pool)
130            .await
131            .map_err(|e| Error::Store(e.to_string()))?;
132        Ok(row.get::<i64, _>("id"))
133    }
134
135    async fn claim_batch(
136        &self,
137        owner_id: &OwnerId,
138        max: usize,
139        visibility: Duration,
140    ) -> Result<Vec<ClaimedBufferEntry<Self::Id>>> {
141        let claimed_until = (Utc::now()
142            + chrono::Duration::from_std(visibility)
143                .map_err(|_| Error::Config("visibility duration out of range".to_owned()))?)
144        .format("%Y-%m-%dT%H:%M:%S%.6fZ")
145        .to_string();
146        let max_i64 = max as i64;
147
148        let mut tx = self
149            .pool
150            .begin()
151            .await
152            .map_err(|e| Error::Store(e.to_string()))?;
153
154        let sql = format!(
155            "WITH picked AS ( \
156                SELECT id FROM {relation} \
157                WHERE claimed_by IS NULL OR claimed_until IS NULL OR claimed_until < NOW() \
158                ORDER BY id \
159                LIMIT $1 \
160                FOR UPDATE SKIP LOCKED \
161            ) \
162            UPDATE {relation} c \
163            SET claimed_by = $2, \
164                claimed_until = $3::timestamptz, \
165                attempts = attempts + 1 \
166            FROM picked \
167            WHERE c.id = picked.id \
168            RETURNING c.id, c.event, c.attempts",
169            relation = self.relation
170        );
171
172        let rows = sqlx::query(&sql)
173            .bind(max_i64)
174            .bind(owner_id.as_str())
175            .bind(claimed_until)
176            .fetch_all(&mut *tx)
177            .await
178            .map_err(|e| Error::Store(e.to_string()))?;
179
180        tx.commit().await.map_err(|e| Error::Store(e.to_string()))?;
181
182        let mut out = Vec::with_capacity(rows.len());
183        for row in rows {
184            let id: i64 = row.get("id");
185            let event_value: serde_json::Value = row.get("event");
186            let attempts: i32 = row.get("attempts");
187            out.push(ClaimedBufferEntry {
188                id,
189                event: decode_event(event_value)?,
190                attempts: attempts as u32,
191            });
192        }
193
194        out.sort_by_key(|e| e.id);
195
196        Ok(out)
197    }
198
199    async fn ack(&self, id: &Self::Id) -> Result<()> {
200        let sql = format!(
201            "DELETE FROM {relation} WHERE id = $1",
202            relation = self.relation
203        );
204        sqlx::query(&sql)
205            .bind(id)
206            .execute(&self.pool)
207            .await
208            .map_err(|e| Error::Store(e.to_string()))?;
209        Ok(())
210    }
211
212    async fn nack(&self, id: &Self::Id) -> Result<()> {
213        let sql = format!(
214            "UPDATE {relation} SET claimed_by = NULL, claimed_until = NULL WHERE id = $1",
215            relation = self.relation
216        );
217        sqlx::query(&sql)
218            .bind(id)
219            .execute(&self.pool)
220            .await
221            .map_err(|e| Error::Store(e.to_string()))?;
222        Ok(())
223    }
224}
225
226#[cfg(test)]
227mod schema_tests {
228    use super::*;
229
230    #[test]
231    fn schema_sql_contains_expected_table() {
232        let sql = PgClaimedBufferStore::schema_sql(&PgClaimedBufferStoreConfig::default());
233        assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"event_buffer_claims\""));
234    }
235}
eventuary_postgres/claim_buffer.rs

eventuary_postgres/
claim_buffer.rs
