Skip to main content

datumesh_store/
pg.rs

1//! `PgStore`: the store on PostgreSQL.
2
3use std::borrow::Cow;
4
5use datumesh_core::lexical::{check_dns_name, check_function_name, check_hex, check_instant};
6use datumesh_core::{
7    json, keys_from_response, sha256_hex, verify, Cell, Descriptor, ErrorCode, Keys, Receipt, Resolver, Scalar,
8};
9use postgres::NoTls;
10
11use crate::error::{Error, Result};
12use crate::pool::{Conn, Pool, PoolOptions};
13use crate::sql::Q;
14use crate::{Namespace, Store, Version};
15
16/// The store on PostgreSQL, with a connection pool. `Send + Sync`: share one per process
17/// (behind an `Arc`); each call checks a connection out for its duration.
18///
19/// Every `put_*` is one transaction. [`Store`] methods return errors: a malformed argument
20/// is [`Error::Invalid`] (no query is made), a database failure [`Error::Database`]. The
21/// [`Resolver`] methods answer both with `None`; a database failure is logged with
22/// `tracing` (target `datumesh_store`, level `error`).
23pub struct PgStore {
24    pool: Pool,
25}
26
27impl std::fmt::Debug for PgStore {
28    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29        f.debug_struct("PgStore").finish_non_exhaustive()
30    }
31}
32
33fn is_hash(s: &str) -> bool {
34    check_hex(s, 64, "$").is_ok()
35}
36
37fn is_instant(s: &str) -> bool {
38    check_instant(s, "$").is_ok()
39}
40
41impl PgStore {
42    /// Connects to a `postgres://` URL or `key=value` connection string, without TLS, with
43    /// the default [`PoolOptions`]. One connection is opened at once, so a wrong URL fails
44    /// here. For TLS, use [`with_connector`](Self::with_connector).
45    pub fn connect(url: &str) -> Result<PgStore> {
46        let config: postgres::Config = url.parse()?;
47        PgStore::with_connector(move || config.connect(NoTls), PoolOptions::new())
48    }
49
50    /// A store whose pool opens connections with `connect` (any TLS setup, any
51    /// configuration). One connection is opened at once, so a failing connector fails here.
52    pub fn with_connector<F>(connect: F, options: PoolOptions) -> Result<PgStore>
53    where
54        F: Fn() -> core::result::Result<postgres::Client, postgres::Error> + Send + Sync + 'static,
55    {
56        let store = PgStore { pool: Pool::new(Box::new(connect), options) };
57        drop(store.pool.get()?);
58        Ok(store)
59    }
60
61    /// Prepares every SQL statement the store uses against the live schema, on a fresh
62    /// connection's behalf (no cached statement is reused), and returns how many there are.
63    /// Fails with the first statement the schema cannot serve.
64    pub fn check_statements(&self) -> Result<usize> {
65        let mut c = self.pool.get()?;
66        for q in Q::ALL {
67            c.client.prepare_typed(q.sql(), q.types())?;
68        }
69        Ok(Q::ALL.len())
70    }
71
72    /// A [`Resolver`] lookup: a failure (logged by `run`) becomes `None`.
73    fn resolve<T>(&self, op: &'static str, f: impl FnOnce(&mut Conn) -> Result<Option<T>>) -> Option<T> {
74        self.run(op, f).ok().flatten()
75    }
76
77    /// Runs `f` on a pooled connection; a database or pool failure is logged.
78    fn run<T>(&self, op: &'static str, f: impl FnOnce(&mut Conn) -> Result<T>) -> Result<T> {
79        let result = self.pool.get().and_then(|mut c| f(&mut c));
80        if let Err(e @ (Error::Database(_) | Error::PoolTimeout)) = &result {
81            tracing::error!(target: "datumesh_store", op, error = %e, "database failure");
82        }
83        result
84    }
85
86    fn version_rows(rows: &[postgres::Row], cell_id: &str) -> Result<Vec<Version>> {
87        rows.iter()
88            .map(|r| Ok(Version { cell_id: cell_id.to_string(), tx: r.try_get(0)?, receipt_hash: r.try_get(1)? }))
89            .collect()
90    }
91}
92
93fn get_bytes(c: &mut Conn, q: Q, key: &str) -> Result<Option<Vec<u8>>> {
94    let st = c.stmt(q)?;
95    let row = c.client.query_opt(&st, &[&key])?;
96    Ok(match row {
97        Some(r) => Some(r.try_get(0)?),
98        None => None,
99    })
100}
101
102fn get_strings(c: &mut Conn, q: Q, key: &str) -> Result<Vec<String>> {
103    let st = c.stmt(q)?;
104    let rows = c.client.query(&st, &[&key])?;
105    rows.iter().map(|r| Ok(r.try_get(0)?)).collect()
106}
107
108impl Resolver for PgStore {
109    fn descriptor(&self, fn_hash: &str) -> Option<Cow<'_, [u8]>> {
110        if !is_hash(fn_hash) {
111            return None;
112        }
113        self.resolve("descriptor", |c| get_bytes(c, Q::DescriptorGet, fn_hash)).map(Cow::Owned)
114    }
115
116    fn receipt(&self, hash: &str) -> Option<Cow<'_, [u8]>> {
117        if !is_hash(hash) {
118            return None;
119        }
120        self.resolve("receipt", |c| get_bytes(c, Q::ReceiptGet, hash)).map(Cow::Owned)
121    }
122
123    /// `Some` only for a `stored` function (descriptor `code.lang` is `stored`, and the store
124    /// has the descriptor): the value of the receipt of the cell's latest version with
125    /// `tx <= asof`. `None` for every other function, and when no version is that old.
126    fn evaluate(&self, cell: &Cell, asof: &str) -> Option<Scalar> {
127        if !is_instant(asof) {
128            return None;
129        }
130        self.resolve("evaluate", |c| {
131            let Some(desc) = get_bytes(c, Q::DescriptorGet, &cell.fn_hash)? else { return Ok(None) };
132            match Descriptor::from_json(&desc) {
133                Ok(d) if d.lang == "stored" => {}
134                Ok(_) => return Ok(None),
135                Err(e) => {
136                    tracing::error!(target: "datumesh_store", fn_hash = %cell.fn_hash, error = %e, "stored descriptor does not parse");
137                    return Ok(None);
138                }
139            }
140            let st = c.stmt(Q::ValueAsof)?;
141            let Some(row) = c.client.query_opt(&st, &[&cell.id(), &asof])? else { return Ok(None) };
142            let bytes: Vec<u8> = row.try_get(0)?;
143            match Receipt::from_json(&bytes) {
144                Ok(r) => Ok(Some(r.value().clone())),
145                Err(e) => {
146                    tracing::error!(target: "datumesh_store", cell_id = %cell.id(), error = %e, "stored receipt does not parse");
147                    Ok(None)
148                }
149            }
150        })
151    }
152}
153
154impl Store for PgStore {
155    fn put_descriptor(&self, bytes: &[u8]) -> Result<String> {
156        let d = Descriptor::from_json(bytes)?;
157        let hash = d.hash();
158        let canonical = d.canonical();
159        self.run("put_descriptor", |c| {
160            let st = c.stmt(Q::DescriptorPut)?;
161            c.client.execute(&st, &[&hash, &canonical])?;
162            Ok(())
163        })?;
164        tracing::debug!(target: "datumesh_store", fn_hash = %hash, "put_descriptor");
165        Ok(hash)
166    }
167
168    fn put_receipt(&self, bytes: &[u8]) -> Result<String> {
169        // Steps 1-4 first: a receipt that is not one is refused without a query.
170        let r = Receipt::from_json(bytes)?;
171        // Steps 1-4 passed, so the text parses; its JCS is the canonical form.
172        let canonical = json::parse(bytes)?.to_jcs();
173        let hash = r.hash().to_string();
174        let cell_id = r.cell().id();
175        let host = r.signer_host().to_string();
176        let inputs: Vec<String> = r.inputs().to_vec();
177        self.run("put_receipt", |c| {
178            let ns = c.stmt(Q::NamespaceGet)?;
179            let put = c.stmt(Q::ReceiptPut)?;
180            let get = c.stmt(Q::ReceiptGet)?;
181            let edges = c.stmt(Q::EdgesPut)?;
182            let mut t = c.client.transaction()?;
183            // Idempotent: the very receipt already stored verified when it was stored, and
184            // the stored copy is unchanged, so it is not checked again (even if its key has
185            // since been revoked).
186            if let Some(row) = t.query_opt(&get, &[&hash])? {
187                let stored: Vec<u8> = row.try_get(0)?;
188                if stored == canonical {
189                    return Ok(());
190                }
191            }
192            // Steps 5-6 against the stored key document of the signer host; none stored is
193            // a host without keys (E_KEY_UNAVAILABLE).
194            let mut keys = Keys::new();
195            if let Some(row) = t.query_opt(&ns, &[&host])? {
196                let doc: Vec<u8> = row.try_get(0)?;
197                keys.insert_response(&host, 200, Some(&doc));
198            }
199            let out = verify(bytes, &keys, None, None);
200            if !out.ok {
201                return Err(Error::Invalid(out.error().unwrap_or_else(|| {
202                    datumesh_core::Error::new(ErrorCode::Signature, "$", "the receipt did not verify")
203                })));
204            }
205            let n = t.execute(&put, &[&hash, &canonical, &cell_id, &r.tx(), &r.cell().fn_hash, &r.publisher()])?;
206            if n == 0 {
207                let row = t.query_one(&get, &[&hash])?;
208                let stored: Vec<u8> = row.try_get(0)?;
209                if stored != canonical {
210                    return Err(Error::Conflict { hash: hash.clone() });
211                }
212            }
213            if !inputs.is_empty() {
214                t.execute(&edges, &[&hash, &inputs])?;
215            }
216            t.commit()?;
217            Ok(())
218        })?;
219        tracing::debug!(target: "datumesh_store", hash = %hash, cell_id = %cell_id, "put_receipt");
220        Ok(hash)
221    }
222
223    fn put_blob(&self, bytes: &[u8]) -> Result<String> {
224        let sha = sha256_hex(bytes);
225        let length = i64::try_from(bytes.len()).unwrap_or(i64::MAX);
226        self.run("put_blob", |c| {
227            let st = c.stmt(Q::BlobPut)?;
228            c.client.execute(&st, &[&sha, &bytes, &length])?;
229            Ok(())
230        })?;
231        tracing::debug!(target: "datumesh_store", sha256 = %sha, length, "put_blob");
232        Ok(sha)
233    }
234
235    fn blob(&self, sha256: &str) -> Result<Option<Cow<'_, [u8]>>> {
236        check_hex(sha256, 64, "sha256")?;
237        Ok(self.run("blob", |c| get_bytes(c, Q::BlobGet, sha256))?.map(Cow::Owned))
238    }
239
240    fn put_version(&self, cell_id: &str, tx: &str, receipt_hash: &str) -> Result<()> {
241        check_hex(cell_id, 64, "cell_id")?;
242        check_instant(tx, "tx")?;
243        check_hex(receipt_hash, 64, "receipt_hash")?;
244        self.run("put_version", |c| {
245            let lock = c.stmt(Q::Lock)?;
246            let meta = c.stmt(Q::ReceiptMeta)?;
247            let latest = c.stmt(Q::VersionLatest)?;
248            let put = c.stmt(Q::VersionPut)?;
249            let mut t = c.client.transaction()?;
250            t.execute(&lock, &[&"versions", &cell_id])?;
251            let Some(row) = t.query_opt(&meta, &[&receipt_hash])? else {
252                return Err(Error::NotFound { what: "receipt", key: receipt_hash.to_string() });
253            };
254            let (r_cell, r_tx): (String, String) = (row.try_get(0)?, row.try_get(1)?);
255            if r_cell != cell_id {
256                return Err(Error::Mismatch { what: "cell_id", expected: cell_id.to_string(), found: r_cell });
257            }
258            if r_tx != tx {
259                return Err(Error::Mismatch { what: "tx", expected: tx.to_string(), found: r_tx });
260            }
261            if let Some(row) = t.query_opt(&latest, &[&cell_id])? {
262                let l: String = row.try_get(0)?;
263                if tx <= l.as_str() {
264                    return Err(Error::NotLater {
265                        what: "version",
266                        key: cell_id.to_string(),
267                        given: tx.to_string(),
268                        latest: l,
269                    });
270                }
271            }
272            t.execute(&put, &[&cell_id, &tx, &receipt_hash])?;
273            t.commit()?;
274            Ok(())
275        })?;
276        tracing::debug!(target: "datumesh_store", cell_id, tx, receipt_hash, "put_version");
277        Ok(())
278    }
279
280    fn version_asof(&self, cell_id: &str, asof: &str) -> Result<Option<Version>> {
281        check_hex(cell_id, 64, "cell_id")?;
282        check_instant(asof, "asof")?;
283        self.run("version_asof", |c| {
284            let st = c.stmt(Q::VersionAsof)?;
285            let rows = c.client.query(&st, &[&cell_id, &asof])?;
286            Ok(PgStore::version_rows(&rows, cell_id)?.pop())
287        })
288    }
289
290    fn versions(&self, cell_id: &str, from: &str, to: &str) -> Result<Vec<Version>> {
291        check_hex(cell_id, 64, "cell_id")?;
292        check_instant(from, "from")?;
293        check_instant(to, "to")?;
294        self.run("versions", |c| {
295            let st = c.stmt(Q::VersionRange)?;
296            let rows = c.client.query(&st, &[&cell_id, &from, &to])?;
297            PgStore::version_rows(&rows, cell_id)
298        })
299    }
300
301    fn latest(&self, cell_id: &str) -> Result<Option<Version>> {
302        check_hex(cell_id, 64, "cell_id")?;
303        self.run("latest", |c| {
304            let st = c.stmt(Q::VersionLatest)?;
305            let rows = c.client.query(&st, &[&cell_id])?;
306            Ok(PgStore::version_rows(&rows, cell_id)?.pop())
307        })
308    }
309
310    fn dependents(&self, receipt_hash: &str) -> Result<Vec<String>> {
311        check_hex(receipt_hash, 64, "receipt_hash")?;
312        self.run("dependents", |c| get_strings(c, Q::Dependents, receipt_hash))
313    }
314
315    fn inputs(&self, receipt_hash: &str) -> Result<Vec<String>> {
316        check_hex(receipt_hash, 64, "receipt_hash")?;
317        self.run("inputs", |c| get_strings(c, Q::Inputs, receipt_hash))
318    }
319
320    fn put_namespace(&self, host: &str, keydoc_bytes: &[u8], validated_at: &str) -> Result<()> {
321        // SPEC 9.3: a 200 with this body, for this host (a DNS name equal to the document's).
322        keys_from_response(host, 200, Some(keydoc_bytes))?;
323        check_instant(validated_at, "validated_at")?;
324        let canonical = json::parse(keydoc_bytes)?.to_jcs();
325        self.run("put_namespace", |c| {
326            let lock = c.stmt(Q::Lock)?;
327            let get = c.stmt(Q::NamespaceGet)?;
328            let put = c.stmt(Q::NamespacePut)?;
329            let mut t = c.client.transaction()?;
330            t.execute(&lock, &[&"namespaces", &host])?;
331            if let Some(row) = t.query_opt(&get, &[&host])? {
332                let at: String = row.try_get(1)?;
333                if validated_at < at.as_str() {
334                    return Err(Error::NotLater {
335                        what: "namespace",
336                        key: host.to_string(),
337                        given: validated_at.to_string(),
338                        latest: at,
339                    });
340                }
341            }
342            t.execute(&put, &[&host, &canonical, &validated_at])?;
343            t.commit()?;
344            Ok(())
345        })?;
346        tracing::debug!(target: "datumesh_store", host, validated_at, "put_namespace");
347        Ok(())
348    }
349
350    fn namespace(&self, host: &str) -> Result<Option<Namespace>> {
351        check_dns_name(host, "host")?;
352        self.run("namespace", |c| {
353            let st = c.stmt(Q::NamespaceGet)?;
354            let Some(row) = c.client.query_opt(&st, &[&host])? else { return Ok(None) };
355            Ok(Some(Namespace {
356                host: host.to_string(),
357                keydoc: row.try_get(0)?,
358                validated_at: row.try_get(1)?,
359                suspended_at: row.try_get(2)?,
360            }))
361        })
362    }
363
364    fn bind_name(&self, fn_name: &str, fn_hash: &str, tx: &str) -> Result<()> {
365        check_function_name(fn_name, "fn_name")?;
366        check_hex(fn_hash, 64, "fn_hash")?;
367        check_instant(tx, "tx")?;
368        self.run("bind_name", |c| {
369            let lock = c.stmt(Q::Lock)?;
370            let exists = c.stmt(Q::DescriptorExists)?;
371            let latest = c.stmt(Q::NameLatest)?;
372            let put = c.stmt(Q::NamePut)?;
373            let mut t = c.client.transaction()?;
374            t.execute(&lock, &[&"names", &fn_name])?;
375            let found: bool = t.query_one(&exists, &[&fn_hash])?.try_get(0)?;
376            if !found {
377                return Err(Error::NotFound { what: "descriptor", key: fn_hash.to_string() });
378            }
379            if let Some(row) = t.query_opt(&latest, &[&fn_name])? {
380                let l: String = row.try_get(0)?;
381                if tx <= l.as_str() {
382                    return Err(Error::NotLater {
383                        what: "name",
384                        key: fn_name.to_string(),
385                        given: tx.to_string(),
386                        latest: l,
387                    });
388                }
389            }
390            t.execute(&put, &[&fn_name, &tx, &fn_hash])?;
391            t.commit()?;
392            Ok(())
393        })?;
394        tracing::debug!(target: "datumesh_store", fn_name, fn_hash, tx, "bind_name");
395        Ok(())
396    }
397
398    fn resolve_name(&self, fn_name: &str, asof: &str) -> Result<Option<String>> {
399        check_function_name(fn_name, "fn_name")?;
400        check_instant(asof, "asof")?;
401        self.run("resolve_name", |c| {
402            let st = c.stmt(Q::NameAsof)?;
403            let row = c.client.query_opt(&st, &[&fn_name, &asof])?;
404            Ok(match row {
405                Some(r) => Some(r.try_get(0)?),
406                None => None,
407            })
408        })
409    }
410}