Skip to main content

eventuary_postgres/
watermark.rs

1//! PostgreSQL [`WatermarkStore`] implementation.
2//!
3//! Persists per-key high-water timestamps. `save_watermark` upserts so
4//! redelivered or out-of-order saves converge.
5
6use std::sync::Arc;
7
8use chrono::{DateTime, Utc};
9use sqlx::{PgPool, Row};
10
11use eventuary_core::io::reader::WatermarkStore;
12use eventuary_core::{Error, Result};
13
14use crate::relation::PgRelationName;
15use crate::schema::{Migration, RelationReplacement};
16
17const WATERMARK_STORE_0001_INIT_SQL: &str = r#"
18CREATE TABLE IF NOT EXISTS {watermarks} (
19    key        TEXT        NOT NULL PRIMARY KEY,
20    ts         TIMESTAMPTZ NOT NULL,
21    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
22);
23"#;
24
25const WATERMARK_STORE_MIGRATIONS: &[Migration] = &[Migration {
26    name: "0001_init",
27    sql: WATERMARK_STORE_0001_INIT_SQL,
28}];
29
30#[derive(Debug, Clone)]
31pub struct PgWatermarkStoreConfig {
32    pub relation: PgRelationName,
33}
34
35impl Default for PgWatermarkStoreConfig {
36    fn default() -> Self {
37        Self {
38            relation: PgRelationName::new("watermarks").expect("default watermarks relation"),
39        }
40    }
41}
42
43#[derive(Clone)]
44pub struct PgWatermarkStore {
45    pool: PgPool,
46    relation: Arc<String>,
47}
48
49impl PgWatermarkStore {
50    pub fn new(pool: PgPool, config: PgWatermarkStoreConfig) -> Self {
51        Self {
52            pool,
53            relation: Arc::new(config.relation.render()),
54        }
55    }
56
57    pub async fn connect(pool: PgPool, config: PgWatermarkStoreConfig) -> Result<Self> {
58        Self::prepare_schema(&pool, &config).await?;
59        Ok(Self::new(pool, config))
60    }
61
62    pub async fn prepare_schema(pool: &PgPool, config: &PgWatermarkStoreConfig) -> Result<()> {
63        crate::schema::apply_schema(
64            pool,
65            WATERMARK_STORE_MIGRATIONS,
66            &[RelationReplacement {
67                token: "{watermarks}",
68                relation: &config.relation,
69            }],
70        )
71        .await
72    }
73
74    pub fn schema_sql(config: &PgWatermarkStoreConfig) -> String {
75        crate::schema::render_schema_sql(
76            WATERMARK_STORE_MIGRATIONS,
77            &[RelationReplacement {
78                token: "{watermarks}",
79                relation: &config.relation,
80            }],
81        )
82    }
83}
84
85impl WatermarkStore for PgWatermarkStore {
86    async fn load_watermark(&self, key: &str) -> Result<Option<DateTime<Utc>>> {
87        let sql = format!(
88            "SELECT to_char(ts AT TIME ZONE 'UTC', 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"') AS ts \
89             FROM {relation} WHERE key = $1",
90            relation = self.relation
91        );
92        let row = sqlx::query(&sql)
93            .bind(key)
94            .fetch_optional(&self.pool)
95            .await
96            .map_err(|e| Error::Store(e.to_string()))?;
97        match row {
98            Some(r) => {
99                let ts_str: String = r.get("ts");
100                let ts = DateTime::parse_from_rfc3339(&ts_str)
101                    .map_err(|e| Error::Serialization(format!("watermark decode: {e}")))?
102                    .with_timezone(&Utc);
103                Ok(Some(ts))
104            }
105            None => Ok(None),
106        }
107    }
108
109    async fn save_watermark(&self, key: &str, ts: DateTime<Utc>) -> Result<()> {
110        let sql = format!(
111            "INSERT INTO {relation} (key, ts) VALUES ($1, $2::timestamptz) \
112             ON CONFLICT (key) DO UPDATE SET ts = EXCLUDED.ts, updated_at = NOW()",
113            relation = self.relation
114        );
115        sqlx::query(&sql)
116            .bind(key)
117            .bind(ts.to_rfc3339())
118            .execute(&self.pool)
119            .await
120            .map_err(|e| Error::Store(e.to_string()))?;
121        Ok(())
122    }
123}
124
125#[cfg(test)]
126mod schema_tests {
127    use super::*;
128
129    #[test]
130    fn schema_sql_contains_expected_table() {
131        let sql = PgWatermarkStore::schema_sql(&PgWatermarkStoreConfig::default());
132        assert!(sql.contains("CREATE TABLE IF NOT EXISTS \"watermarks\""));
133    }
134}