eventuary_postgres/
watermark_store.rs1use 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}