prolly_store_postgres/
lib.rs

1#![doc = include_str!("../README.md")]
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 std::collections::{BTreeMap, BTreeSet};
11    use std::num::NonZeroUsize;
12
13    use sqlx::{PgConnection, PgPool, Row};
14
15    use crate::{
16        RemoteBatchOp, RemoteManifestUpdate, RemoteNamedRoot, RemoteRootCondition, RemoteRootWrite,
17        RemoteStoreBackend, RemoteTransactionConflict, RemoteTransactionUpdate,
18    };
19
20    /// Store adapter for PostgreSQL-backed prolly nodes and roots.
21    pub type PostgresStore = crate::RemoteProllyStore<PostgresBackend>;
22
23    /// PostgreSQL adapter tuning that does not change stored data.
24    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
25    pub struct PostgresBackendOptions {
26        max_batch_items: NonZeroUsize,
27    }
28
29    impl PostgresBackendOptions {
30        /// Create options with the maximum number of items sent in one SQL batch.
31        pub const fn new(max_batch_items: NonZeroUsize) -> Self {
32            Self { max_batch_items }
33        }
34
35        /// Maximum number of items sent in one SQL batch.
36        pub const fn max_batch_items(self) -> usize {
37            self.max_batch_items.get()
38        }
39    }
40
41    impl Default for PostgresBackendOptions {
42        fn default() -> Self {
43            Self::new(NonZeroUsize::new(1_024).expect("1024 is nonzero"))
44        }
45    }
46
47    /// SQLx-backed PostgreSQL backend.
48    #[derive(Clone, Debug)]
49    pub struct PostgresBackend {
50        pool: PgPool,
51        options: PostgresBackendOptions,
52    }
53
54    impl PostgresBackend {
55        /// Create a backend from an existing SQLx pool.
56        pub fn new(pool: PgPool) -> Self {
57            Self::new_with_options(pool, PostgresBackendOptions::default())
58        }
59
60        /// Create a backend from an existing SQLx pool and adapter options.
61        pub fn new_with_options(pool: PgPool, options: PostgresBackendOptions) -> Self {
62            Self { pool, options }
63        }
64
65        /// Connect to PostgreSQL using `database_url`.
66        pub async fn connect(database_url: &str) -> Result<Self, sqlx::Error> {
67            Self::connect_with_options(database_url, PostgresBackendOptions::default()).await
68        }
69
70        /// Connect to PostgreSQL using `database_url` and adapter options.
71        pub async fn connect_with_options(
72            database_url: &str,
73            options: PostgresBackendOptions,
74        ) -> Result<Self, sqlx::Error> {
75            Ok(Self::new_with_options(
76                PgPool::connect(database_url).await?,
77                options,
78            ))
79        }
80
81        /// Borrow the underlying pool.
82        pub fn pool(&self) -> &PgPool {
83            &self.pool
84        }
85
86        /// Return this backend's adapter options.
87        pub const fn options(&self) -> PostgresBackendOptions {
88            self.options
89        }
90
91        /// Create the required tables if they do not already exist.
92        pub async fn initialize_schema(&self) -> Result<(), sqlx::Error> {
93            execute_statements(&self.pool, POSTGRES_SCHEMA).await
94        }
95    }
96
97    impl RemoteStoreBackend for PostgresBackend {
98        type Error = sqlx::Error;
99
100        async fn get_node(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
101            sqlx::query("SELECT node FROM prolly_nodes WHERE cid = $1")
102                .bind(key)
103                .fetch_optional(&self.pool)
104                .await?
105                .map(|row| row.try_get("node"))
106                .transpose()
107        }
108
109        async fn put_node(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
110            sqlx::query(
111                "\
112                INSERT INTO prolly_nodes (cid, node) VALUES ($1, $2) \
113                ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
114            )
115            .bind(key)
116            .bind(value)
117            .execute(&self.pool)
118            .await?;
119            Ok(())
120        }
121
122        async fn delete_node(&self, key: &[u8]) -> Result<(), Self::Error> {
123            sqlx::query("DELETE FROM prolly_nodes WHERE cid = $1")
124                .bind(key)
125                .execute(&self.pool)
126                .await?;
127            Ok(())
128        }
129
130        async fn batch_nodes(&self, ops: &[RemoteBatchOp<'_>]) -> Result<(), Self::Error> {
131            if ops.is_empty() {
132                return Ok(());
133            }
134            let mut final_ops = BTreeMap::<Vec<u8>, Option<Vec<u8>>>::new();
135            for op in ops {
136                match op {
137                    RemoteBatchOp::Upsert { key, value } => {
138                        final_ops.insert((*key).to_vec(), Some((*value).to_vec()));
139                    }
140                    RemoteBatchOp::Delete { key } => {
141                        final_ops.insert((*key).to_vec(), None);
142                    }
143                }
144            }
145            let deletes = final_ops
146                .iter()
147                .filter_map(|(key, value)| value.is_none().then_some(key.as_slice()))
148                .collect::<Vec<_>>();
149            let upserts = final_ops
150                .iter()
151                .filter_map(|(key, value)| value.as_deref().map(|value| (key.as_slice(), value)))
152                .collect::<Vec<_>>();
153            let mut tx = self.pool.begin().await?;
154            delete_node_chunks(&mut tx, &deletes, self.options.max_batch_items()).await?;
155            upsert_node_chunks(&mut tx, &upserts, self.options.max_batch_items()).await?;
156            tx.commit().await
157        }
158
159        async fn batch_get_nodes_ordered(
160            &self,
161            keys: &[&[u8]],
162        ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
163            if keys.is_empty() {
164                return Ok(Vec::new());
165            }
166            let mut values = Vec::with_capacity(keys.len());
167            for chunk in keys.chunks(self.options.max_batch_items()) {
168                let requested = chunk.iter().map(|key| (*key).to_vec()).collect::<Vec<_>>();
169                let rows = sqlx::query(
170                    "\
171                    SELECT requested.ord, nodes.node \
172                    FROM unnest($1::bytea[]) WITH ORDINALITY AS requested(cid, ord) \
173                    LEFT JOIN prolly_nodes AS nodes ON nodes.cid = requested.cid \
174                    ORDER BY requested.ord",
175                )
176                .bind(requested)
177                .fetch_all(&self.pool)
178                .await?;
179                for row in rows {
180                    values.push(row.try_get::<Option<Vec<u8>>, _>("node")?);
181                }
182            }
183            Ok(values)
184        }
185
186        async fn batch_put_nodes(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
187            if entries.is_empty() {
188                return Ok(());
189            }
190            let entries = deduplicate_entries(entries);
191            let entries = entries
192                .iter()
193                .map(|(key, value)| (key.as_slice(), value.as_slice()))
194                .collect::<Vec<_>>();
195            let mut tx = self.pool.begin().await?;
196            upsert_node_chunks(&mut tx, &entries, self.options.max_batch_items()).await?;
197            tx.commit().await
198        }
199
200        async fn list_node_cids(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
201            let rows = sqlx::query("SELECT cid FROM prolly_nodes ORDER BY cid")
202                .fetch_all(&self.pool)
203                .await?;
204            rows.into_iter().map(|row| row.try_get("cid")).collect()
205        }
206
207        fn prefers_batch_reads(&self) -> bool {
208            true
209        }
210
211        fn supports_hints(&self) -> bool {
212            true
213        }
214
215        async fn get_hint(
216            &self,
217            namespace: &[u8],
218            key: &[u8],
219        ) -> Result<Option<Vec<u8>>, Self::Error> {
220            sqlx::query("SELECT value FROM prolly_hints WHERE namespace = $1 AND key = $2")
221                .bind(namespace)
222                .bind(key)
223                .fetch_optional(&self.pool)
224                .await?
225                .map(|row| row.try_get("value"))
226                .transpose()
227        }
228
229        async fn put_hint(
230            &self,
231            namespace: &[u8],
232            key: &[u8],
233            value: &[u8],
234        ) -> Result<(), Self::Error> {
235            sqlx::query(
236                "\
237                INSERT INTO prolly_hints (namespace, key, value) VALUES ($1, $2, $3) \
238                ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
239            )
240            .bind(namespace)
241            .bind(key)
242            .bind(value)
243            .execute(&self.pool)
244            .await?;
245            Ok(())
246        }
247
248        async fn batch_put_nodes_with_hint(
249            &self,
250            entries: &[(&[u8], &[u8])],
251            namespace: &[u8],
252            key: &[u8],
253            value: &[u8],
254        ) -> Result<(), Self::Error> {
255            let entries = deduplicate_entries(entries);
256            let entries = entries
257                .iter()
258                .map(|(key, value)| (key.as_slice(), value.as_slice()))
259                .collect::<Vec<_>>();
260            let mut tx = self.pool.begin().await?;
261            upsert_node_chunks(&mut tx, &entries, self.options.max_batch_items()).await?;
262            sqlx::query(
263                "\
264                INSERT INTO prolly_hints (namespace, key, value) VALUES ($1, $2, $3) \
265                ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
266            )
267            .bind(namespace)
268            .bind(key)
269            .bind(value)
270            .execute(&mut *tx)
271            .await?;
272            tx.commit().await
273        }
274
275        async fn get_root_manifest(&self, name: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
276            sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
277                .bind(name)
278                .fetch_optional(&self.pool)
279                .await?
280                .map(|row| row.try_get("manifest"))
281                .transpose()
282        }
283
284        async fn put_root_manifest(&self, name: &[u8], manifest: &[u8]) -> Result<(), Self::Error> {
285            let mut tx = self.pool.begin().await?;
286            lock_root_names(&mut tx, &[name.to_vec()]).await?;
287            upsert_root_chunks(&mut tx, &[(name, manifest)], self.options.max_batch_items())
288                .await?;
289            tx.commit().await
290        }
291
292        async fn delete_root_manifest(&self, name: &[u8]) -> Result<(), Self::Error> {
293            let mut tx = self.pool.begin().await?;
294            lock_root_names(&mut tx, &[name.to_vec()]).await?;
295            delete_root_chunks(&mut tx, &[name], self.options.max_batch_items()).await?;
296            tx.commit().await
297        }
298
299        async fn compare_and_swap_root_manifest(
300            &self,
301            name: &[u8],
302            expected: Option<&[u8]>,
303            new: Option<&[u8]>,
304        ) -> Result<RemoteManifestUpdate, Self::Error> {
305            let mut tx = self.pool.begin().await?;
306            lock_root_names(&mut tx, &[name.to_vec()]).await?;
307
308            let current = sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
309                .bind(name)
310                .fetch_optional(&mut *tx)
311                .await?
312                .map(|row| row.try_get("manifest"))
313                .transpose()?;
314            if current.as_deref() != expected {
315                tx.rollback().await?;
316                return Ok(RemoteManifestUpdate::Conflict { current });
317            }
318
319            match new {
320                Some(manifest) => {
321                    upsert_root_chunks(
322                        &mut tx,
323                        &[(name, manifest)],
324                        self.options.max_batch_items(),
325                    )
326                    .await?;
327                }
328                None => {
329                    delete_root_chunks(&mut tx, &[name], self.options.max_batch_items()).await?;
330                }
331            }
332
333            tx.commit().await?;
334            Ok(RemoteManifestUpdate::Applied)
335        }
336
337        async fn list_root_manifests(&self) -> Result<Vec<RemoteNamedRoot>, Self::Error> {
338            let rows = sqlx::query("SELECT name, manifest FROM prolly_roots ORDER BY name")
339                .fetch_all(&self.pool)
340                .await?;
341            rows.into_iter()
342                .map(|row| {
343                    Ok(RemoteNamedRoot::new(
344                        row.try_get("name")?,
345                        row.try_get("manifest")?,
346                    ))
347                })
348                .collect()
349        }
350
351        fn supports_transactions(&self) -> bool {
352            true
353        }
354
355        async fn commit_transaction(
356            &self,
357            node_writes: &[RemoteBatchOp<'_>],
358            root_conditions: &[RemoteRootCondition],
359            root_writes: &[RemoteRootWrite],
360        ) -> Result<RemoteTransactionUpdate, Self::Error> {
361            let mut tx = self.pool.begin().await?;
362            let root_names = root_names(root_conditions, root_writes);
363            lock_root_names(&mut tx, &root_names).await?;
364            let current_roots = read_root_manifests(&mut tx, &root_names).await?;
365
366            for condition in root_conditions {
367                let current = current_roots.get(&condition.name).cloned().unwrap_or(None);
368                if current != condition.expected {
369                    tx.rollback().await?;
370                    return Ok(RemoteTransactionUpdate::Conflict(
371                        RemoteTransactionConflict::new(
372                            condition.name.clone(),
373                            condition.expected.clone(),
374                            current,
375                        ),
376                    ));
377                }
378            }
379
380            let mut final_nodes = BTreeMap::<Vec<u8>, Option<Vec<u8>>>::new();
381            for write in node_writes {
382                match write {
383                    RemoteBatchOp::Upsert { key, value } => {
384                        final_nodes.insert((*key).to_vec(), Some((*value).to_vec()));
385                    }
386                    RemoteBatchOp::Delete { key } => {
387                        final_nodes.insert((*key).to_vec(), None);
388                    }
389                };
390            }
391            let node_deletes = final_nodes
392                .iter()
393                .filter_map(|(key, value)| value.is_none().then_some(key.as_slice()))
394                .collect::<Vec<_>>();
395            let node_upserts = final_nodes
396                .iter()
397                .filter_map(|(key, value)| value.as_deref().map(|value| (key.as_slice(), value)))
398                .collect::<Vec<_>>();
399            delete_node_chunks(&mut tx, &node_deletes, self.options.max_batch_items()).await?;
400            upsert_node_chunks(&mut tx, &node_upserts, self.options.max_batch_items()).await?;
401
402            let mut final_roots = BTreeMap::<Vec<u8>, Option<Vec<u8>>>::new();
403            for write in root_writes {
404                match write {
405                    RemoteRootWrite::Put { name, manifest } => {
406                        final_roots.insert(name.clone(), Some(manifest.clone()));
407                    }
408                    RemoteRootWrite::Delete { name } => {
409                        final_roots.insert(name.clone(), None);
410                    }
411                }
412            }
413            let root_deletes = final_roots
414                .iter()
415                .filter_map(|(name, manifest)| manifest.is_none().then_some(name.as_slice()))
416                .collect::<Vec<_>>();
417            let root_upserts = final_roots
418                .iter()
419                .filter_map(|(name, manifest)| {
420                    manifest
421                        .as_deref()
422                        .map(|manifest| (name.as_slice(), manifest))
423                })
424                .collect::<Vec<_>>();
425            delete_root_chunks(&mut tx, &root_deletes, self.options.max_batch_items()).await?;
426            upsert_root_chunks(&mut tx, &root_upserts, self.options.max_batch_items()).await?;
427
428            tx.commit().await?;
429            Ok(RemoteTransactionUpdate::Applied)
430        }
431    }
432
433    async fn upsert_node_chunks(
434        connection: &mut PgConnection,
435        entries: &[(&[u8], &[u8])],
436        max_batch_items: usize,
437    ) -> Result<(), sqlx::Error> {
438        for chunk in entries.chunks(max_batch_items) {
439            let keys = chunk
440                .iter()
441                .map(|(key, _)| (*key).to_vec())
442                .collect::<Vec<_>>();
443            let values = chunk
444                .iter()
445                .map(|(_, value)| (*value).to_vec())
446                .collect::<Vec<_>>();
447            sqlx::query(
448                "\
449                INSERT INTO prolly_nodes (cid, node) \
450                SELECT input.cid, input.node \
451                FROM unnest($1::bytea[], $2::bytea[]) AS input(cid, node) \
452                ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
453            )
454            .bind(keys)
455            .bind(values)
456            .execute(&mut *connection)
457            .await?;
458        }
459        Ok(())
460    }
461
462    fn deduplicate_entries(entries: &[(&[u8], &[u8])]) -> BTreeMap<Vec<u8>, Vec<u8>> {
463        entries
464            .iter()
465            .map(|(key, value)| ((*key).to_vec(), (*value).to_vec()))
466            .collect()
467    }
468
469    async fn delete_node_chunks(
470        connection: &mut PgConnection,
471        keys: &[&[u8]],
472        max_batch_items: usize,
473    ) -> Result<(), sqlx::Error> {
474        for chunk in keys.chunks(max_batch_items) {
475            let keys = chunk.iter().map(|key| (*key).to_vec()).collect::<Vec<_>>();
476            sqlx::query("DELETE FROM prolly_nodes WHERE cid = ANY($1::bytea[])")
477                .bind(keys)
478                .execute(&mut *connection)
479                .await?;
480        }
481        Ok(())
482    }
483
484    async fn upsert_root_chunks(
485        connection: &mut PgConnection,
486        entries: &[(&[u8], &[u8])],
487        max_batch_items: usize,
488    ) -> Result<(), sqlx::Error> {
489        for chunk in entries.chunks(max_batch_items) {
490            let names = chunk
491                .iter()
492                .map(|(name, _)| (*name).to_vec())
493                .collect::<Vec<_>>();
494            let manifests = chunk
495                .iter()
496                .map(|(_, manifest)| (*manifest).to_vec())
497                .collect::<Vec<_>>();
498            sqlx::query(
499                "\
500                INSERT INTO prolly_roots (name, manifest) \
501                SELECT input.name, input.manifest \
502                FROM unnest($1::bytea[], $2::bytea[]) AS input(name, manifest) \
503                ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest",
504            )
505            .bind(names)
506            .bind(manifests)
507            .execute(&mut *connection)
508            .await?;
509        }
510        Ok(())
511    }
512
513    async fn delete_root_chunks(
514        connection: &mut PgConnection,
515        names: &[&[u8]],
516        max_batch_items: usize,
517    ) -> Result<(), sqlx::Error> {
518        for chunk in names.chunks(max_batch_items) {
519            let names = chunk
520                .iter()
521                .map(|name| (*name).to_vec())
522                .collect::<Vec<_>>();
523            sqlx::query("DELETE FROM prolly_roots WHERE name = ANY($1::bytea[])")
524                .bind(names)
525                .execute(&mut *connection)
526                .await?;
527        }
528        Ok(())
529    }
530
531    fn root_names(conditions: &[RemoteRootCondition], writes: &[RemoteRootWrite]) -> Vec<Vec<u8>> {
532        let mut names = BTreeSet::new();
533        names.extend(conditions.iter().map(|condition| condition.name.clone()));
534        names.extend(writes.iter().map(|write| match write {
535            RemoteRootWrite::Put { name, .. } | RemoteRootWrite::Delete { name } => name.clone(),
536        }));
537        names.into_iter().collect()
538    }
539
540    async fn lock_root_names(
541        connection: &mut PgConnection,
542        names: &[Vec<u8>],
543    ) -> Result<(), sqlx::Error> {
544        for name in names {
545            sqlx::query(
546                "\
547                SELECT pg_advisory_xact_lock( \
548                    hashtextextended('prolly-root-v1:' || encode($1::bytea, 'hex'), 0) \
549                )",
550            )
551            .bind(name)
552            .execute(&mut *connection)
553            .await?;
554        }
555        Ok(())
556    }
557
558    async fn read_root_manifests(
559        connection: &mut PgConnection,
560        names: &[Vec<u8>],
561    ) -> Result<BTreeMap<Vec<u8>, Option<Vec<u8>>>, sqlx::Error> {
562        if names.is_empty() {
563            return Ok(BTreeMap::new());
564        }
565        let rows = sqlx::query(
566            "\
567            SELECT requested.name, roots.manifest \
568            FROM unnest($1::bytea[]) AS requested(name) \
569            LEFT JOIN prolly_roots AS roots ON roots.name = requested.name",
570        )
571        .bind(names)
572        .fetch_all(&mut *connection)
573        .await?;
574        rows.into_iter()
575            .map(|row| {
576                Ok((
577                    row.try_get::<Vec<u8>, _>("name")?,
578                    row.try_get::<Option<Vec<u8>>, _>("manifest")?,
579                ))
580            })
581            .collect()
582    }
583
584    async fn execute_statements(pool: &PgPool, sql: &str) -> Result<(), sqlx::Error> {
585        for statement in sql
586            .split(';')
587            .map(str::trim)
588            .filter(|stmt| !stmt.is_empty())
589        {
590            sqlx::query(statement).execute(pool).await?;
591        }
592        Ok(())
593    }
594
595    /// Minimal table layout for PostgreSQL implementations.
596    pub const POSTGRES_SCHEMA: &str = "\
597CREATE TABLE IF NOT EXISTS prolly_nodes (
598  cid bytea PRIMARY KEY,
599  node bytea NOT NULL
600);
601CREATE TABLE IF NOT EXISTS prolly_hints (
602  namespace bytea NOT NULL,
603  key bytea NOT NULL,
604  value bytea NOT NULL,
605  PRIMARY KEY(namespace, key)
606);
607CREATE TABLE IF NOT EXISTS prolly_roots (
608  name bytea PRIMARY KEY,
609  manifest bytea NOT NULL
610);";
611}
612
613pub use postgres::*;
prolly_store_postgres/lib.rs

prolly_store_postgres/
lib.rs
