prolly_store_spanner/
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/// Spanner adapter entry point.
9pub mod spanner {
10    use google_cloud_googleapis::spanner::v1::Mutation;
11    use google_cloud_spanner::client::{Client, ClientConfig, Error};
12    use google_cloud_spanner::key::Key;
13    use google_cloud_spanner::mutation::{delete, insert_or_update};
14    use google_cloud_spanner::statement::Statement;
15    use google_cloud_spanner::transaction_rw::ReadWriteTransaction;
16
17    use crate::{
18        RemoteBatchOp, RemoteManifestUpdate, RemoteNamedRoot, RemoteRootCondition, RemoteRootWrite,
19        RemoteStoreBackend, RemoteTransactionConflict, RemoteTransactionUpdate,
20    };
21
22    /// Store adapter for Spanner-backed prolly nodes and roots.
23    pub type SpannerStore = crate::RemoteProllyStore<SpannerBackend>;
24
25    /// Google Cloud Spanner-backed backend.
26    #[derive(Clone)]
27    pub struct SpannerBackend {
28        client: Client,
29        read_parallelism: usize,
30    }
31
32    impl std::fmt::Debug for SpannerBackend {
33        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34            f.debug_struct("SpannerBackend")
35                .field("read_parallelism", &self.read_parallelism)
36                .finish_non_exhaustive()
37        }
38    }
39
40    impl SpannerBackend {
41        /// Create a backend from an existing Spanner client.
42        pub fn new(client: Client) -> Self {
43            Self {
44                client,
45                read_parallelism: DEFAULT_READ_PARALLELISM,
46            }
47        }
48
49        /// Connect to a Spanner database resource name using a caller-provided config.
50        pub async fn connect(database: &str, config: ClientConfig) -> Result<Self, Error> {
51            Ok(Self::new(Client::new(database, config).await?))
52        }
53
54        /// Borrow the underlying Spanner client.
55        pub fn client(&self) -> &Client {
56            &self.client
57        }
58
59        /// Set the read parallelism advertised to async prolly traversals.
60        pub fn with_read_parallelism(mut self, read_parallelism: usize) -> Self {
61            self.read_parallelism = read_parallelism.max(1);
62            self
63        }
64
65        async fn query_one_value(
66            &self,
67            statement: Statement,
68            column: &str,
69        ) -> Result<Option<Vec<u8>>, Error> {
70            let mut tx = self.client.single().await?;
71            let mut rows = tx.query(statement).await.map_err(Error::from)?;
72            let row = rows.next().await.map_err(Error::from)?;
73            row.map(|row| row.column_by_name(column).map_err(Error::from))
74                .transpose()
75        }
76
77        async fn query_bytes_column(
78            &self,
79            statement: Statement,
80            column: &str,
81        ) -> Result<Vec<Vec<u8>>, Error> {
82            let mut tx = self.client.single().await?;
83            let mut rows = tx.query(statement).await.map_err(Error::from)?;
84            let mut values = Vec::new();
85            while let Some(row) = rows.next().await.map_err(Error::from)? {
86                values.push(row.column_by_name(column)?);
87            }
88            Ok(values)
89        }
90    }
91
92    impl RemoteStoreBackend for SpannerBackend {
93        type Error = Error;
94
95        async fn get_node(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
96            let mut statement = Statement::new("SELECT Node FROM ProllyNodes WHERE Cid = @cid");
97            statement.add_param("cid", &key.to_vec());
98            self.query_one_value(statement, "Node").await
99        }
100
101        async fn put_node(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
102            self.client
103                .apply(vec![node_upsert(key, value)])
104                .await
105                .map(|_| ())
106        }
107
108        async fn delete_node(&self, key: &[u8]) -> Result<(), Self::Error> {
109            self.client.apply(vec![node_delete(key)]).await.map(|_| ())
110        }
111
112        async fn batch_nodes(&self, ops: &[RemoteBatchOp<'_>]) -> Result<(), Self::Error> {
113            let mutations = ops
114                .iter()
115                .map(|op| match op {
116                    RemoteBatchOp::Upsert { key, value } => node_upsert(key, value),
117                    RemoteBatchOp::Delete { key } => node_delete(key),
118                })
119                .collect::<Vec<_>>();
120            if mutations.is_empty() {
121                return Ok(());
122            }
123            self.client.apply(mutations).await.map(|_| ())
124        }
125
126        async fn batch_get_nodes_ordered(
127            &self,
128            keys: &[&[u8]],
129        ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
130            let mut values = Vec::with_capacity(keys.len());
131            for key in keys {
132                values.push(self.get_node(key).await?);
133            }
134            Ok(values)
135        }
136
137        async fn batch_put_nodes(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
138            let mutations = entries
139                .iter()
140                .map(|(key, value)| node_upsert(key, value))
141                .collect::<Vec<_>>();
142            if mutations.is_empty() {
143                return Ok(());
144            }
145            self.client.apply(mutations).await.map(|_| ())
146        }
147
148        async fn list_node_cids(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
149            self.query_bytes_column(
150                Statement::new("SELECT Cid FROM ProllyNodes ORDER BY Cid"),
151                "Cid",
152            )
153            .await
154        }
155
156        fn read_parallelism(&self) -> usize {
157            self.read_parallelism
158        }
159
160        fn supports_hints(&self) -> bool {
161            true
162        }
163
164        async fn get_hint(
165            &self,
166            namespace: &[u8],
167            key: &[u8],
168        ) -> Result<Option<Vec<u8>>, Self::Error> {
169            let mut statement = Statement::new(
170                "SELECT Value FROM ProllyHints WHERE Namespace = @namespace AND HintKey = @key",
171            );
172            statement.add_param("namespace", &namespace.to_vec());
173            statement.add_param("key", &key.to_vec());
174            self.query_one_value(statement, "Value").await
175        }
176
177        async fn put_hint(
178            &self,
179            namespace: &[u8],
180            key: &[u8],
181            value: &[u8],
182        ) -> Result<(), Self::Error> {
183            self.client
184                .apply(vec![hint_upsert(namespace, key, value)])
185                .await
186                .map(|_| ())
187        }
188
189        async fn batch_put_nodes_with_hint(
190            &self,
191            entries: &[(&[u8], &[u8])],
192            namespace: &[u8],
193            key: &[u8],
194            value: &[u8],
195        ) -> Result<(), Self::Error> {
196            let mut mutations = entries
197                .iter()
198                .map(|(key, value)| node_upsert(key, value))
199                .collect::<Vec<_>>();
200            mutations.push(hint_upsert(namespace, key, value));
201            self.client.apply(mutations).await.map(|_| ())
202        }
203
204        async fn get_root_manifest(&self, name: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
205            let mut statement =
206                Statement::new("SELECT Manifest FROM ProllyRoots WHERE Name = @name");
207            statement.add_param("name", &name.to_vec());
208            self.query_one_value(statement, "Manifest").await
209        }
210
211        async fn put_root_manifest(&self, name: &[u8], manifest: &[u8]) -> Result<(), Self::Error> {
212            self.client
213                .apply(vec![root_upsert(name, manifest)])
214                .await
215                .map(|_| ())
216        }
217
218        async fn delete_root_manifest(&self, name: &[u8]) -> Result<(), Self::Error> {
219            self.client.apply(vec![root_delete(name)]).await.map(|_| ())
220        }
221
222        async fn compare_and_swap_root_manifest(
223            &self,
224            name: &[u8],
225            expected: Option<&[u8]>,
226            new: Option<&[u8]>,
227        ) -> Result<RemoteManifestUpdate, Self::Error> {
228            let name = name.to_vec();
229            let expected = expected.map(<[u8]>::to_vec);
230            let new = new.map(<[u8]>::to_vec);
231            let (_, update) = self
232                .client
233                .read_write_transaction(|tx| {
234                    let name = name.clone();
235                    let expected = expected.clone();
236                    let new = new.clone();
237                    Box::pin(async move {
238                        let current = read_root_in_transaction(tx, &name).await?;
239                        if current.as_deref() != expected.as_deref() {
240                            return Ok::<RemoteManifestUpdate, Error>(
241                                RemoteManifestUpdate::Conflict { current },
242                            );
243                        }
244
245                        match new {
246                            Some(manifest) => tx.buffer_write(vec![root_upsert(&name, &manifest)]),
247                            None => tx.buffer_write(vec![root_delete(&name)]),
248                        }
249                        Ok::<RemoteManifestUpdate, Error>(RemoteManifestUpdate::Applied)
250                    })
251                })
252                .await?;
253            Ok(update)
254        }
255
256        async fn list_root_manifests(&self) -> Result<Vec<RemoteNamedRoot>, Self::Error> {
257            let mut tx = self.client.single().await?;
258            let mut rows = tx
259                .query(Statement::new(
260                    "SELECT Name, Manifest FROM ProllyRoots ORDER BY Name",
261                ))
262                .await
263                .map_err(Error::from)?;
264            let mut roots = Vec::new();
265            while let Some(row) = rows.next().await.map_err(Error::from)? {
266                roots.push(RemoteNamedRoot::new(
267                    row.column_by_name("Name")?,
268                    row.column_by_name("Manifest")?,
269                ));
270            }
271            Ok(roots)
272        }
273
274        fn supports_transactions(&self) -> bool {
275            true
276        }
277
278        async fn commit_transaction(
279            &self,
280            node_writes: &[RemoteBatchOp<'_>],
281            root_conditions: &[RemoteRootCondition],
282            root_writes: &[RemoteRootWrite],
283        ) -> Result<RemoteTransactionUpdate, Self::Error> {
284            let node_writes = node_writes
285                .iter()
286                .map(|write| match write {
287                    RemoteBatchOp::Upsert { key, value } => (true, key.to_vec(), value.to_vec()),
288                    RemoteBatchOp::Delete { key } => (false, key.to_vec(), Vec::new()),
289                })
290                .collect::<Vec<_>>();
291            let root_conditions = root_conditions.to_vec();
292            let root_writes = root_writes.to_vec();
293
294            let (_, update) = self
295                .client
296                .read_write_transaction(|tx| {
297                    let node_writes = node_writes.clone();
298                    let root_conditions = root_conditions.clone();
299                    let root_writes = root_writes.clone();
300                    Box::pin(async move {
301                        for condition in &root_conditions {
302                            let current = read_root_in_transaction(tx, &condition.name).await?;
303                            if current != condition.expected {
304                                return Ok::<RemoteTransactionUpdate, Error>(
305                                    RemoteTransactionUpdate::Conflict(
306                                        RemoteTransactionConflict::new(
307                                            condition.name.clone(),
308                                            condition.expected.clone(),
309                                            current,
310                                        ),
311                                    ),
312                                );
313                            }
314                        }
315
316                        let mut mutations = Vec::new();
317                        for (is_upsert, key, value) in &node_writes {
318                            if *is_upsert {
319                                mutations.push(node_upsert(key, value));
320                            } else {
321                                mutations.push(node_delete(key));
322                            }
323                        }
324                        for write in &root_writes {
325                            match write {
326                                RemoteRootWrite::Put { name, manifest } => {
327                                    mutations.push(root_upsert(name, manifest));
328                                }
329                                RemoteRootWrite::Delete { name } => {
330                                    mutations.push(root_delete(name));
331                                }
332                            }
333                        }
334                        if !mutations.is_empty() {
335                            tx.buffer_write(mutations);
336                        }
337                        Ok::<RemoteTransactionUpdate, Error>(RemoteTransactionUpdate::Applied)
338                    })
339                })
340                .await?;
341            Ok(update)
342        }
343    }
344
345    async fn read_root_in_transaction(
346        tx: &mut ReadWriteTransaction,
347        name: &[u8],
348    ) -> Result<Option<Vec<u8>>, Error> {
349        let name = name.to_vec();
350        let row = tx
351            .read_row(ROOTS_TABLE, &["Manifest"], Key::new(&name))
352            .await
353            .map_err(Error::from)?;
354        row.map(|row| row.column_by_name("Manifest").map_err(Error::from))
355            .transpose()
356    }
357
358    fn node_upsert(key: &[u8], value: &[u8]) -> Mutation {
359        let key = key.to_vec();
360        let value = value.to_vec();
361        insert_or_update(NODES_TABLE, &["Cid", "Node"], &[&key, &value])
362    }
363
364    fn node_delete(key: &[u8]) -> Mutation {
365        let key = key.to_vec();
366        delete(NODES_TABLE, Key::new(&key))
367    }
368
369    fn hint_upsert(namespace: &[u8], key: &[u8], value: &[u8]) -> Mutation {
370        let namespace = namespace.to_vec();
371        let key = key.to_vec();
372        let value = value.to_vec();
373        insert_or_update(
374            HINTS_TABLE,
375            &["Namespace", "HintKey", "Value"],
376            &[&namespace, &key, &value],
377        )
378    }
379
380    fn root_upsert(name: &[u8], manifest: &[u8]) -> Mutation {
381        let name = name.to_vec();
382        let manifest = manifest.to_vec();
383        insert_or_update(ROOTS_TABLE, &["Name", "Manifest"], &[&name, &manifest])
384    }
385
386    fn root_delete(name: &[u8]) -> Mutation {
387        let name = name.to_vec();
388        delete(ROOTS_TABLE, Key::new(&name))
389    }
390
391    const DEFAULT_READ_PARALLELISM: usize = 16;
392    const NODES_TABLE: &str = "ProllyNodes";
393    const HINTS_TABLE: &str = "ProllyHints";
394    const ROOTS_TABLE: &str = "ProllyRoots";
395
396    /// Minimal GoogleSQL table layout for Spanner implementations.
397    pub const SPANNER_SCHEMA: &str = "\
398CREATE TABLE ProllyNodes (
399  Cid BYTES(32) NOT NULL,
400  Node BYTES(MAX) NOT NULL
401) PRIMARY KEY (Cid);
402CREATE TABLE ProllyHints (
403  Namespace BYTES(MAX) NOT NULL,
404  HintKey BYTES(MAX) NOT NULL,
405  Value BYTES(MAX) NOT NULL
406) PRIMARY KEY (Namespace, HintKey);
407CREATE TABLE ProllyRoots (
408  Name BYTES(MAX) NOT NULL,
409  Manifest BYTES(MAX) NOT NULL
410) PRIMARY KEY (Name);";
411}
412
413pub use spanner::*;
prolly_store_spanner/lib.rs

prolly_store_spanner/
lib.rs
