prolly_store_postgres/
lib.rs

1//! PostgreSQL store adapter for prolly-map.
2
3pub use prolly::{
4    RemoteBatchOp, RemoteManifestUpdate, RemoteNamedRoot, RemoteProllyStore, RemoteRootCondition,
5    RemoteRootWrite, RemoteStoreBackend, RemoteTransactionConflict, RemoteTransactionUpdate,
6};
7
8/// Postgres adapter entry point.
9pub mod postgres {
10    use sqlx::{PgPool, Row};
11
12    use crate::{
13        RemoteBatchOp, RemoteManifestUpdate, RemoteNamedRoot, RemoteRootCondition, RemoteRootWrite,
14        RemoteStoreBackend, RemoteTransactionConflict, RemoteTransactionUpdate,
15    };
16
17    /// Store adapter for PostgreSQL-backed prolly nodes and roots.
18    pub type PostgresStore = crate::RemoteProllyStore<PostgresBackend>;
19
20    /// SQLx-backed PostgreSQL backend.
21    #[derive(Clone, Debug)]
22    pub struct PostgresBackend {
23        pool: PgPool,
24    }
25
26    impl PostgresBackend {
27        /// Create a backend from an existing SQLx pool.
28        pub fn new(pool: PgPool) -> Self {
29            Self { pool }
30        }
31
32        /// Connect to PostgreSQL using `database_url`.
33        pub async fn connect(database_url: &str) -> Result<Self, sqlx::Error> {
34            Ok(Self::new(PgPool::connect(database_url).await?))
35        }
36
37        /// Borrow the underlying pool.
38        pub fn pool(&self) -> &PgPool {
39            &self.pool
40        }
41
42        /// Create the required tables if they do not already exist.
43        pub async fn initialize_schema(&self) -> Result<(), sqlx::Error> {
44            execute_statements(&self.pool, POSTGRES_SCHEMA).await
45        }
46    }
47
48    impl RemoteStoreBackend for PostgresBackend {
49        type Error = sqlx::Error;
50
51        async fn get_node(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
52            sqlx::query("SELECT node FROM prolly_nodes WHERE cid = $1")
53                .bind(key)
54                .fetch_optional(&self.pool)
55                .await?
56                .map(|row| row.try_get("node"))
57                .transpose()
58        }
59
60        async fn put_node(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
61            sqlx::query(
62                "\
63                INSERT INTO prolly_nodes (cid, node) VALUES ($1, $2) \
64                ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
65            )
66            .bind(key)
67            .bind(value)
68            .execute(&self.pool)
69            .await?;
70            Ok(())
71        }
72
73        async fn delete_node(&self, key: &[u8]) -> Result<(), Self::Error> {
74            sqlx::query("DELETE FROM prolly_nodes WHERE cid = $1")
75                .bind(key)
76                .execute(&self.pool)
77                .await?;
78            Ok(())
79        }
80
81        async fn batch_nodes(&self, ops: &[RemoteBatchOp<'_>]) -> Result<(), Self::Error> {
82            let mut tx = self.pool.begin().await?;
83            for op in ops {
84                match op {
85                    RemoteBatchOp::Upsert { key, value } => {
86                        sqlx::query(
87                            "\
88                            INSERT INTO prolly_nodes (cid, node) VALUES ($1, $2) \
89                            ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
90                        )
91                        .bind(*key)
92                        .bind(*value)
93                        .execute(&mut *tx)
94                        .await?;
95                    }
96                    RemoteBatchOp::Delete { key } => {
97                        sqlx::query("DELETE FROM prolly_nodes WHERE cid = $1")
98                            .bind(*key)
99                            .execute(&mut *tx)
100                            .await?;
101                    }
102                }
103            }
104            tx.commit().await
105        }
106
107        async fn batch_get_nodes_ordered(
108            &self,
109            keys: &[&[u8]],
110        ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
111            let mut values = Vec::with_capacity(keys.len());
112            for key in keys {
113                values.push(self.get_node(key).await?);
114            }
115            Ok(values)
116        }
117
118        async fn batch_put_nodes(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
119            let mut tx = self.pool.begin().await?;
120            for (key, value) in entries {
121                sqlx::query(
122                    "\
123                    INSERT INTO prolly_nodes (cid, node) VALUES ($1, $2) \
124                    ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
125                )
126                .bind(*key)
127                .bind(*value)
128                .execute(&mut *tx)
129                .await?;
130            }
131            tx.commit().await
132        }
133
134        async fn list_node_cids(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
135            let rows = sqlx::query("SELECT cid FROM prolly_nodes ORDER BY cid")
136                .fetch_all(&self.pool)
137                .await?;
138            rows.into_iter().map(|row| row.try_get("cid")).collect()
139        }
140
141        fn prefers_batch_reads(&self) -> bool {
142            true
143        }
144
145        fn supports_hints(&self) -> bool {
146            true
147        }
148
149        async fn get_hint(
150            &self,
151            namespace: &[u8],
152            key: &[u8],
153        ) -> Result<Option<Vec<u8>>, Self::Error> {
154            sqlx::query("SELECT value FROM prolly_hints WHERE namespace = $1 AND key = $2")
155                .bind(namespace)
156                .bind(key)
157                .fetch_optional(&self.pool)
158                .await?
159                .map(|row| row.try_get("value"))
160                .transpose()
161        }
162
163        async fn put_hint(
164            &self,
165            namespace: &[u8],
166            key: &[u8],
167            value: &[u8],
168        ) -> Result<(), Self::Error> {
169            sqlx::query(
170                "\
171                INSERT INTO prolly_hints (namespace, key, value) VALUES ($1, $2, $3) \
172                ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
173            )
174            .bind(namespace)
175            .bind(key)
176            .bind(value)
177            .execute(&self.pool)
178            .await?;
179            Ok(())
180        }
181
182        async fn batch_put_nodes_with_hint(
183            &self,
184            entries: &[(&[u8], &[u8])],
185            namespace: &[u8],
186            key: &[u8],
187            value: &[u8],
188        ) -> Result<(), Self::Error> {
189            let mut tx = self.pool.begin().await?;
190            for (key, value) in entries {
191                sqlx::query(
192                    "\
193                    INSERT INTO prolly_nodes (cid, node) VALUES ($1, $2) \
194                    ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
195                )
196                .bind(*key)
197                .bind(*value)
198                .execute(&mut *tx)
199                .await?;
200            }
201            sqlx::query(
202                "\
203                INSERT INTO prolly_hints (namespace, key, value) VALUES ($1, $2, $3) \
204                ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
205            )
206            .bind(namespace)
207            .bind(key)
208            .bind(value)
209            .execute(&mut *tx)
210            .await?;
211            tx.commit().await
212        }
213
214        async fn get_root_manifest(&self, name: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
215            sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
216                .bind(name)
217                .fetch_optional(&self.pool)
218                .await?
219                .map(|row| row.try_get("manifest"))
220                .transpose()
221        }
222
223        async fn put_root_manifest(&self, name: &[u8], manifest: &[u8]) -> Result<(), Self::Error> {
224            sqlx::query(
225                "\
226                INSERT INTO prolly_roots (name, manifest) VALUES ($1, $2) \
227                ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest",
228            )
229            .bind(name)
230            .bind(manifest)
231            .execute(&self.pool)
232            .await?;
233            Ok(())
234        }
235
236        async fn delete_root_manifest(&self, name: &[u8]) -> Result<(), Self::Error> {
237            sqlx::query("DELETE FROM prolly_roots WHERE name = $1")
238                .bind(name)
239                .execute(&self.pool)
240                .await?;
241            Ok(())
242        }
243
244        async fn compare_and_swap_root_manifest(
245            &self,
246            name: &[u8],
247            expected: Option<&[u8]>,
248            new: Option<&[u8]>,
249        ) -> Result<RemoteManifestUpdate, Self::Error> {
250            let mut tx = self.pool.begin().await?;
251            sqlx::query("LOCK TABLE prolly_roots IN SHARE ROW EXCLUSIVE MODE")
252                .execute(&mut *tx)
253                .await?;
254
255            let current = sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
256                .bind(name)
257                .fetch_optional(&mut *tx)
258                .await?
259                .map(|row| row.try_get("manifest"))
260                .transpose()?;
261            if current.as_deref() != expected {
262                tx.rollback().await?;
263                return Ok(RemoteManifestUpdate::Conflict { current });
264            }
265
266            match new {
267                Some(manifest) => {
268                    sqlx::query(
269                        "\
270                        INSERT INTO prolly_roots (name, manifest) VALUES ($1, $2) \
271                        ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest",
272                    )
273                    .bind(name)
274                    .bind(manifest)
275                    .execute(&mut *tx)
276                    .await?;
277                }
278                None => {
279                    sqlx::query("DELETE FROM prolly_roots WHERE name = $1")
280                        .bind(name)
281                        .execute(&mut *tx)
282                        .await?;
283                }
284            }
285
286            tx.commit().await?;
287            Ok(RemoteManifestUpdate::Applied)
288        }
289
290        async fn list_root_manifests(&self) -> Result<Vec<RemoteNamedRoot>, Self::Error> {
291            let rows = sqlx::query("SELECT name, manifest FROM prolly_roots ORDER BY name")
292                .fetch_all(&self.pool)
293                .await?;
294            rows.into_iter()
295                .map(|row| {
296                    Ok(RemoteNamedRoot::new(
297                        row.try_get("name")?,
298                        row.try_get("manifest")?,
299                    ))
300                })
301                .collect()
302        }
303
304        fn supports_transactions(&self) -> bool {
305            true
306        }
307
308        async fn commit_transaction(
309            &self,
310            node_writes: &[RemoteBatchOp<'_>],
311            root_conditions: &[RemoteRootCondition],
312            root_writes: &[RemoteRootWrite],
313        ) -> Result<RemoteTransactionUpdate, Self::Error> {
314            let mut tx = self.pool.begin().await?;
315            sqlx::query("LOCK TABLE prolly_roots IN SHARE ROW EXCLUSIVE MODE")
316                .execute(&mut *tx)
317                .await?;
318
319            for condition in root_conditions {
320                let current = sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
321                    .bind(&condition.name)
322                    .fetch_optional(&mut *tx)
323                    .await?
324                    .map(|row| row.try_get("manifest"))
325                    .transpose()?;
326                if current != condition.expected {
327                    tx.rollback().await?;
328                    return Ok(RemoteTransactionUpdate::Conflict(
329                        RemoteTransactionConflict::new(
330                            condition.name.clone(),
331                            condition.expected.clone(),
332                            current,
333                        ),
334                    ));
335                }
336            }
337
338            for write in node_writes {
339                match write {
340                    RemoteBatchOp::Upsert { key, value } => {
341                        sqlx::query(
342                            "\
343                            INSERT INTO prolly_nodes (cid, node) VALUES ($1, $2) \
344                            ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
345                        )
346                        .bind(*key)
347                        .bind(*value)
348                        .execute(&mut *tx)
349                        .await?;
350                    }
351                    RemoteBatchOp::Delete { key } => {
352                        sqlx::query("DELETE FROM prolly_nodes WHERE cid = $1")
353                            .bind(*key)
354                            .execute(&mut *tx)
355                            .await?;
356                    }
357                }
358            }
359
360            for write in root_writes {
361                match write {
362                    RemoteRootWrite::Put { name, manifest } => {
363                        sqlx::query(
364                            "\
365                            INSERT INTO prolly_roots (name, manifest) VALUES ($1, $2) \
366                            ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest",
367                        )
368                        .bind(name)
369                        .bind(manifest)
370                        .execute(&mut *tx)
371                        .await?;
372                    }
373                    RemoteRootWrite::Delete { name } => {
374                        sqlx::query("DELETE FROM prolly_roots WHERE name = $1")
375                            .bind(name)
376                            .execute(&mut *tx)
377                            .await?;
378                    }
379                }
380            }
381
382            tx.commit().await?;
383            Ok(RemoteTransactionUpdate::Applied)
384        }
385    }
386
387    async fn execute_statements(pool: &PgPool, sql: &str) -> Result<(), sqlx::Error> {
388        for statement in sql
389            .split(';')
390            .map(str::trim)
391            .filter(|stmt| !stmt.is_empty())
392        {
393            sqlx::query(statement).execute(pool).await?;
394        }
395        Ok(())
396    }
397
398    /// Minimal table layout for PostgreSQL implementations.
399    pub const POSTGRES_SCHEMA: &str = "\
400CREATE TABLE IF NOT EXISTS prolly_nodes (
401  cid bytea PRIMARY KEY,
402  node bytea NOT NULL
403);
404CREATE TABLE IF NOT EXISTS prolly_hints (
405  namespace bytea NOT NULL,
406  key bytea NOT NULL,
407  value bytea NOT NULL,
408  PRIMARY KEY(namespace, key)
409);
410CREATE TABLE IF NOT EXISTS prolly_roots (
411  name bytea PRIMARY KEY,
412  manifest bytea NOT NULL
413);";
414}
415
416pub use postgres::*;
prolly_store_postgres/lib.rs

prolly_store_postgres/
lib.rs
