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, HashMap};
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 = HashMap::<&[u8], Option<&[u8]>>::with_capacity(ops.len());
135            for op in ops {
136                match op {
137                    RemoteBatchOp::Upsert { key, value } => {
138                        final_ops.insert(key, Some(value));
139                    }
140                    RemoteBatchOp::Delete { key } => {
141                        final_ops.insert(key, None);
142                    }
143                }
144            }
145            let deletes = final_ops
146                .iter()
147                .filter_map(|(key, value)| value.is_none().then_some(*key))
148                .collect::<Vec<_>>();
149            let upserts = final_ops
150                .iter()
151                .filter_map(|(key, value)| value.map(|value| (*key, 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.into_iter().collect::<Vec<_>>();
192            let mut tx = self.pool.begin().await?;
193            upsert_node_chunks(&mut tx, &entries, self.options.max_batch_items()).await?;
194            tx.commit().await
195        }
196
197        async fn list_node_cids(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
198            let rows = sqlx::query("SELECT cid FROM prolly_nodes ORDER BY cid")
199                .fetch_all(&self.pool)
200                .await?;
201            rows.into_iter().map(|row| row.try_get("cid")).collect()
202        }
203
204        fn prefers_batch_reads(&self) -> bool {
205            true
206        }
207
208        fn supports_hints(&self) -> bool {
209            true
210        }
211
212        async fn get_hint(
213            &self,
214            namespace: &[u8],
215            key: &[u8],
216        ) -> Result<Option<Vec<u8>>, Self::Error> {
217            sqlx::query("SELECT value FROM prolly_hints WHERE namespace = $1 AND key = $2")
218                .bind(namespace)
219                .bind(key)
220                .fetch_optional(&self.pool)
221                .await?
222                .map(|row| row.try_get("value"))
223                .transpose()
224        }
225
226        async fn put_hint(
227            &self,
228            namespace: &[u8],
229            key: &[u8],
230            value: &[u8],
231        ) -> Result<(), Self::Error> {
232            sqlx::query(
233                "\
234                INSERT INTO prolly_hints (namespace, key, value) VALUES ($1, $2, $3) \
235                ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
236            )
237            .bind(namespace)
238            .bind(key)
239            .bind(value)
240            .execute(&self.pool)
241            .await?;
242            Ok(())
243        }
244
245        async fn batch_put_nodes_with_hint(
246            &self,
247            entries: &[(&[u8], &[u8])],
248            namespace: &[u8],
249            key: &[u8],
250            value: &[u8],
251        ) -> Result<(), Self::Error> {
252            let entries = deduplicate_entries(entries);
253            let entries = entries.into_iter().collect::<Vec<_>>();
254            let mut tx = self.pool.begin().await?;
255            upsert_node_chunks(&mut tx, &entries, self.options.max_batch_items()).await?;
256            sqlx::query(
257                "\
258                INSERT INTO prolly_hints (namespace, key, value) VALUES ($1, $2, $3) \
259                ON CONFLICT(namespace, key) DO UPDATE SET value = excluded.value",
260            )
261            .bind(namespace)
262            .bind(key)
263            .bind(value)
264            .execute(&mut *tx)
265            .await?;
266            tx.commit().await
267        }
268
269        async fn get_root_manifest(&self, name: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
270            sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
271                .bind(name)
272                .fetch_optional(&self.pool)
273                .await?
274                .map(|row| row.try_get("manifest"))
275                .transpose()
276        }
277
278        async fn put_root_manifest(&self, name: &[u8], manifest: &[u8]) -> Result<(), Self::Error> {
279            let mut tx = self.pool.begin().await?;
280            lock_root_names(&mut tx, &[name.to_vec()]).await?;
281            upsert_root_chunks(&mut tx, &[(name, manifest)], self.options.max_batch_items())
282                .await?;
283            tx.commit().await
284        }
285
286        async fn delete_root_manifest(&self, name: &[u8]) -> Result<(), Self::Error> {
287            let mut tx = self.pool.begin().await?;
288            lock_root_names(&mut tx, &[name.to_vec()]).await?;
289            delete_root_chunks(&mut tx, &[name], self.options.max_batch_items()).await?;
290            tx.commit().await
291        }
292
293        async fn compare_and_swap_root_manifest(
294            &self,
295            name: &[u8],
296            expected: Option<&[u8]>,
297            new: Option<&[u8]>,
298        ) -> Result<RemoteManifestUpdate, Self::Error> {
299            let mut tx = self.pool.begin().await?;
300            lock_root_names(&mut tx, &[name.to_vec()]).await?;
301
302            let current = sqlx::query("SELECT manifest FROM prolly_roots WHERE name = $1")
303                .bind(name)
304                .fetch_optional(&mut *tx)
305                .await?
306                .map(|row| row.try_get("manifest"))
307                .transpose()?;
308            if current.as_deref() != expected {
309                tx.rollback().await?;
310                return Ok(RemoteManifestUpdate::Conflict { current });
311            }
312
313            match new {
314                Some(manifest) => {
315                    upsert_root_chunks(
316                        &mut tx,
317                        &[(name, manifest)],
318                        self.options.max_batch_items(),
319                    )
320                    .await?;
321                }
322                None => {
323                    delete_root_chunks(&mut tx, &[name], self.options.max_batch_items()).await?;
324                }
325            }
326
327            tx.commit().await?;
328            Ok(RemoteManifestUpdate::Applied)
329        }
330
331        async fn list_root_manifests(&self) -> Result<Vec<RemoteNamedRoot>, Self::Error> {
332            let rows = sqlx::query("SELECT name, manifest FROM prolly_roots ORDER BY name")
333                .fetch_all(&self.pool)
334                .await?;
335            rows.into_iter()
336                .map(|row| {
337                    Ok(RemoteNamedRoot::new(
338                        row.try_get("name")?,
339                        row.try_get("manifest")?,
340                    ))
341                })
342                .collect()
343        }
344
345        fn supports_transactions(&self) -> bool {
346            true
347        }
348
349        async fn commit_transaction(
350            &self,
351            node_writes: &[RemoteBatchOp<'_>],
352            root_conditions: &[RemoteRootCondition],
353            root_writes: &[RemoteRootWrite],
354        ) -> Result<RemoteTransactionUpdate, Self::Error> {
355            let mut tx = self.pool.begin().await?;
356            let root_names = root_names(root_conditions, root_writes);
357            lock_root_names(&mut tx, &root_names).await?;
358            let current_roots = read_root_manifests(&mut tx, &root_names).await?;
359
360            for condition in root_conditions {
361                let current = current_roots.get(&condition.name).cloned().unwrap_or(None);
362                if current != condition.expected {
363                    tx.rollback().await?;
364                    return Ok(RemoteTransactionUpdate::Conflict(
365                        RemoteTransactionConflict::new(
366                            condition.name.clone(),
367                            condition.expected.clone(),
368                            current,
369                        ),
370                    ));
371                }
372            }
373
374            let mut final_nodes =
375                HashMap::<&[u8], Option<&[u8]>>::with_capacity(node_writes.len());
376            for write in node_writes {
377                match write {
378                    RemoteBatchOp::Upsert { key, value } => {
379                        final_nodes.insert(key, Some(value));
380                    }
381                    RemoteBatchOp::Delete { key } => {
382                        final_nodes.insert(key, None);
383                    }
384                };
385            }
386            let node_deletes = final_nodes
387                .iter()
388                .filter_map(|(key, value)| value.is_none().then_some(*key))
389                .collect::<Vec<_>>();
390            let node_upserts = final_nodes
391                .iter()
392                .filter_map(|(key, value)| value.map(|value| (*key, value)))
393                .collect::<Vec<_>>();
394            delete_node_chunks(&mut tx, &node_deletes, self.options.max_batch_items()).await?;
395            upsert_node_chunks(&mut tx, &node_upserts, self.options.max_batch_items()).await?;
396
397            let mut final_roots = BTreeMap::<Vec<u8>, Option<Vec<u8>>>::new();
398            for write in root_writes {
399                match write {
400                    RemoteRootWrite::Put { name, manifest } => {
401                        final_roots.insert(name.clone(), Some(manifest.clone()));
402                    }
403                    RemoteRootWrite::Delete { name } => {
404                        final_roots.insert(name.clone(), None);
405                    }
406                }
407            }
408            let root_deletes = final_roots
409                .iter()
410                .filter_map(|(name, manifest)| manifest.is_none().then_some(name.as_slice()))
411                .collect::<Vec<_>>();
412            let root_upserts = final_roots
413                .iter()
414                .filter_map(|(name, manifest)| {
415                    manifest
416                        .as_deref()
417                        .map(|manifest| (name.as_slice(), manifest))
418                })
419                .collect::<Vec<_>>();
420            delete_root_chunks(&mut tx, &root_deletes, self.options.max_batch_items()).await?;
421            upsert_root_chunks(&mut tx, &root_upserts, self.options.max_batch_items()).await?;
422
423            tx.commit().await?;
424            Ok(RemoteTransactionUpdate::Applied)
425        }
426    }
427
428    async fn upsert_node_chunks(
429        connection: &mut PgConnection,
430        entries: &[(&[u8], &[u8])],
431        max_batch_items: usize,
432    ) -> Result<(), sqlx::Error> {
433        for chunk in entries.chunks(max_batch_items) {
434            let keys = chunk
435                .iter()
436                .map(|(key, _)| (*key).to_vec())
437                .collect::<Vec<_>>();
438            let values = chunk
439                .iter()
440                .map(|(_, value)| (*value).to_vec())
441                .collect::<Vec<_>>();
442            sqlx::query(
443                "\
444                INSERT INTO prolly_nodes (cid, node) \
445                SELECT input.cid, input.node \
446                FROM unnest($1::bytea[], $2::bytea[]) AS input(cid, node) \
447                ON CONFLICT(cid) DO UPDATE SET node = excluded.node",
448            )
449            .bind(keys)
450            .bind(values)
451            .execute(&mut *connection)
452            .await?;
453        }
454        Ok(())
455    }
456
457    fn deduplicate_entries<'a>(entries: &[(&'a [u8], &'a [u8])]) -> HashMap<&'a [u8], &'a [u8]> {
458        entries.iter().copied().collect()
459    }
460
461    async fn delete_node_chunks(
462        connection: &mut PgConnection,
463        keys: &[&[u8]],
464        max_batch_items: usize,
465    ) -> Result<(), sqlx::Error> {
466        for chunk in keys.chunks(max_batch_items) {
467            let keys = chunk.iter().map(|key| (*key).to_vec()).collect::<Vec<_>>();
468            sqlx::query("DELETE FROM prolly_nodes WHERE cid = ANY($1::bytea[])")
469                .bind(keys)
470                .execute(&mut *connection)
471                .await?;
472        }
473        Ok(())
474    }
475
476    async fn upsert_root_chunks(
477        connection: &mut PgConnection,
478        entries: &[(&[u8], &[u8])],
479        max_batch_items: usize,
480    ) -> Result<(), sqlx::Error> {
481        for chunk in entries.chunks(max_batch_items) {
482            let names = chunk
483                .iter()
484                .map(|(name, _)| (*name).to_vec())
485                .collect::<Vec<_>>();
486            let manifests = chunk
487                .iter()
488                .map(|(_, manifest)| (*manifest).to_vec())
489                .collect::<Vec<_>>();
490            sqlx::query(
491                "\
492                INSERT INTO prolly_roots (name, manifest) \
493                SELECT input.name, input.manifest \
494                FROM unnest($1::bytea[], $2::bytea[]) AS input(name, manifest) \
495                ON CONFLICT(name) DO UPDATE SET manifest = excluded.manifest",
496            )
497            .bind(names)
498            .bind(manifests)
499            .execute(&mut *connection)
500            .await?;
501        }
502        Ok(())
503    }
504
505    async fn delete_root_chunks(
506        connection: &mut PgConnection,
507        names: &[&[u8]],
508        max_batch_items: usize,
509    ) -> Result<(), sqlx::Error> {
510        for chunk in names.chunks(max_batch_items) {
511            let names = chunk
512                .iter()
513                .map(|name| (*name).to_vec())
514                .collect::<Vec<_>>();
515            sqlx::query("DELETE FROM prolly_roots WHERE name = ANY($1::bytea[])")
516                .bind(names)
517                .execute(&mut *connection)
518                .await?;
519        }
520        Ok(())
521    }
522
523    fn root_names(conditions: &[RemoteRootCondition], writes: &[RemoteRootWrite]) -> Vec<Vec<u8>> {
524        let mut names = BTreeSet::new();
525        names.extend(conditions.iter().map(|condition| condition.name.clone()));
526        names.extend(writes.iter().map(|write| match write {
527            RemoteRootWrite::Put { name, .. } | RemoteRootWrite::Delete { name } => name.clone(),
528        }));
529        names.into_iter().collect()
530    }
531
532    async fn lock_root_names(
533        connection: &mut PgConnection,
534        names: &[Vec<u8>],
535    ) -> Result<(), sqlx::Error> {
536        for name in names {
537            sqlx::query(
538                "\
539                SELECT pg_advisory_xact_lock( \
540                    hashtextextended('prolly-root-v1:' || encode($1::bytea, 'hex'), 0) \
541                )",
542            )
543            .bind(name)
544            .execute(&mut *connection)
545            .await?;
546        }
547        Ok(())
548    }
549
550    async fn read_root_manifests(
551        connection: &mut PgConnection,
552        names: &[Vec<u8>],
553    ) -> Result<BTreeMap<Vec<u8>, Option<Vec<u8>>>, sqlx::Error> {
554        if names.is_empty() {
555            return Ok(BTreeMap::new());
556        }
557        let rows = sqlx::query(
558            "\
559            SELECT requested.name, roots.manifest \
560            FROM unnest($1::bytea[]) AS requested(name) \
561            LEFT JOIN prolly_roots AS roots ON roots.name = requested.name",
562        )
563        .bind(names)
564        .fetch_all(&mut *connection)
565        .await?;
566        rows.into_iter()
567            .map(|row| {
568                Ok((
569                    row.try_get::<Vec<u8>, _>("name")?,
570                    row.try_get::<Option<Vec<u8>>, _>("manifest")?,
571                ))
572            })
573            .collect()
574    }
575
576    async fn execute_statements(pool: &PgPool, sql: &str) -> Result<(), sqlx::Error> {
577        for statement in sql
578            .split(';')
579            .map(str::trim)
580            .filter(|stmt| !stmt.is_empty())
581        {
582            sqlx::query(statement).execute(pool).await?;
583        }
584        Ok(())
585    }
586
587    /// Minimal table layout for PostgreSQL implementations.
588    pub const POSTGRES_SCHEMA: &str = "\
589CREATE TABLE IF NOT EXISTS prolly_nodes (
590  cid bytea PRIMARY KEY,
591  node bytea NOT NULL
592);
593CREATE TABLE IF NOT EXISTS prolly_hints (
594  namespace bytea NOT NULL,
595  key bytea NOT NULL,
596  value bytea NOT NULL,
597  PRIMARY KEY(namespace, key)
598);
599CREATE TABLE IF NOT EXISTS prolly_roots (
600  name bytea PRIMARY KEY,
601  manifest bytea NOT NULL
602);";
603}
604
605pub use postgres::*;
prolly_store_postgres/lib.rs

prolly_store_postgres/
lib.rs
