Skip to main content

nstreams_postgres/
store.rs

1use async_trait::async_trait;
2use nstreams_core::event::{EventVersion, StreamEvent};
3use nstreams_core::namespace::{Namespace, NamespaceResources};
4use nstreams_core::store::EventStore;
5use sqlx::postgres::PgPoolOptions;
6use sqlx::{PgPool, Row};
7use thiserror::Error;
8
9#[derive(Debug, Error)]
10pub enum PostgresError {
11    #[error(transparent)]
12    Sqlx(#[from] sqlx::Error),
13}
14
15impl From<PostgresError> for nstreams_core::Error {
16    fn from(value: PostgresError) -> Self {
17        Self::Store(value.to_string())
18    }
19}
20
21fn store_err(error: sqlx::Error) -> nstreams_core::Error {
22    PostgresError::from(error).into()
23}
24
25fn row_to_event(
26    namespace: &str,
27    version: i64,
28    payload: serde_json::Value,
29    filter_value: Option<String>,
30) -> nstreams_core::Result<StreamEvent> {
31    Ok(StreamEvent {
32        namespace: namespace.to_string(),
33        version: version as EventVersion,
34        payload: serde_json::to_vec(&payload)
35            .map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?,
36        filter_value,
37    })
38}
39
40/// PostgreSQL-backed event store with per-namespace tables.
41#[derive(Clone)]
42pub struct PostgresEventStore {
43    pool: PgPool,
44}
45
46impl PostgresEventStore {
47    pub async fn connect(database_url: &str) -> Result<Self, PostgresError> {
48        let pool = PgPoolOptions::new()
49            .max_connections(10)
50            .connect(database_url)
51            .await?;
52        Ok(Self { pool })
53    }
54
55    pub fn from_pool(pool: PgPool) -> Self {
56        Self { pool }
57    }
58
59    pub fn pool(&self) -> &PgPool {
60        &self.pool
61    }
62}
63
64#[async_trait]
65impl EventStore for PostgresEventStore {
66    async fn ensure_namespace<N: Namespace>(&self, namespace: &N) -> nstreams_core::Result<()> {
67        let resources = NamespaceResources::new(namespace);
68        let table = resources.postgres_table;
69
70        let ddl = format!(
71            r#"
72            CREATE TABLE IF NOT EXISTS {table} (
73                version BIGINT PRIMARY KEY,
74                payload JSONB NOT NULL,
75                filter_value TEXT,
76                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
77            );
78            CREATE INDEX IF NOT EXISTS idx_{table}_version_desc ON {table} (version DESC);
79            ALTER TABLE {table} ADD COLUMN IF NOT EXISTS filter_value TEXT;
80            "#
81        );
82
83        sqlx::raw_sql(&ddl).execute(&self.pool).await.map_err(store_err)?;
84        Ok(())
85    }
86
87    async fn persist_next<N: Namespace>(
88        &self,
89        namespace: &N,
90        payload: &[u8],
91        filter_value: Option<&str>,
92    ) -> nstreams_core::Result<StreamEvent> {
93        let resources = NamespaceResources::new(namespace);
94        let table = resources.postgres_table;
95        let namespace = resources.namespace;
96
97        let json: serde_json::Value = serde_json::from_slice(payload)
98            .map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?;
99
100        let mut tx = self.pool.begin().await.map_err(store_err)?;
101
102        let lock_sql = format!("SELECT pg_advisory_xact_lock(hashtext($1))");
103        sqlx::query(&lock_sql)
104            .bind(&namespace)
105            .execute(&mut *tx)
106            .await
107            .map_err(store_err)?;
108
109        let next_version_sql = format!(
110            "SELECT COALESCE(MAX(version), 0) + 1 AS version FROM {table}"
111        );
112        let row = sqlx::query(&next_version_sql)
113            .fetch_one(&mut *tx)
114            .await
115            .map_err(store_err)?;
116        let version: i64 = row.get("version");
117
118        let insert_sql = format!(
119            "INSERT INTO {table} (version, payload, filter_value) VALUES ($1, $2, $3)"
120        );
121        sqlx::query(&insert_sql)
122            .bind(version)
123            .bind(&json)
124            .bind(filter_value)
125            .execute(&mut *tx)
126            .await
127            .map_err(store_err)?;
128
129        tx.commit().await.map_err(store_err)?;
130
131        Ok(StreamEvent {
132            namespace,
133            version: version as EventVersion,
134            payload: payload.to_vec(),
135            filter_value: filter_value.map(str::to_string),
136        })
137    }
138
139    async fn load_history<N: Namespace>(
140        &self,
141        namespace: &N,
142        limit: u64,
143    ) -> nstreams_core::Result<Vec<StreamEvent>> {
144        let resources = NamespaceResources::new(namespace);
145        let table = resources.postgres_table;
146        let namespace = resources.namespace;
147
148        let sql = format!(
149            r#"
150            SELECT version, payload, filter_value
151            FROM {table}
152            ORDER BY version DESC
153            LIMIT $1
154            "#
155        );
156
157        let rows = sqlx::query(&sql)
158            .bind(limit as i64)
159            .fetch_all(&self.pool)
160            .await
161            .map_err(store_err)?;
162
163        let mut events = rows
164            .into_iter()
165            .map(|row| {
166                let version: i64 = row.get("version");
167                let payload: serde_json::Value = row.get("payload");
168                let filter_value: Option<String> = row.get("filter_value");
169                row_to_event(&namespace, version, payload, filter_value)
170            })
171            .collect::<nstreams_core::Result<Vec<_>>>()?;
172
173        events.sort_by_key(|event| event.version);
174        Ok(events)
175    }
176
177    async fn load_after_version<N: Namespace>(
178        &self,
179        namespace: &N,
180        after_version: EventVersion,
181        limit: u64,
182    ) -> nstreams_core::Result<Vec<StreamEvent>> {
183        let resources = NamespaceResources::new(namespace);
184        let table = resources.postgres_table;
185        let namespace = resources.namespace;
186
187        let sql = format!(
188            r#"
189            SELECT version, payload, filter_value
190            FROM {table}
191            WHERE version > $1
192            ORDER BY version ASC
193            LIMIT $2
194            "#
195        );
196
197        let rows = sqlx::query(&sql)
198            .bind(after_version as i64)
199            .bind(limit as i64)
200            .fetch_all(&self.pool)
201            .await
202            .map_err(store_err)?;
203
204        rows.into_iter()
205            .map(|row| {
206                let version: i64 = row.get("version");
207                let payload: serde_json::Value = row.get("payload");
208                let filter_value: Option<String> = row.get("filter_value");
209                row_to_event(&namespace, version, payload, filter_value)
210            })
211            .collect()
212    }
213
214    async fn latest_version<N: Namespace>(
215        &self,
216        namespace: &N,
217    ) -> nstreams_core::Result<Option<EventVersion>> {
218        let resources = NamespaceResources::new(namespace);
219        let table = resources.postgres_table;
220
221        let sql = format!("SELECT MAX(version) AS version FROM {table}");
222        let row = sqlx::query(&sql)
223            .fetch_optional(&self.pool)
224            .await
225            .map_err(store_err)?;
226        Ok(row.and_then(|row| {
227            let version: Option<i64> = row.get("version");
228            version.map(|v| v as EventVersion)
229        }))
230    }
231}