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