nstreams_postgres/
store.rs1use 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#[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}