Skip to main content

hyperdb_api/
kv_store.rs

1// Copyright (c) 2026, Salesforce, Inc. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! Key-value store over a fixed Hyper table.
5//!
6//! [`KvStore`] is an ergonomic string-native KV abstraction backed by a
7//! single table, `KV_TABLE`, namespaced by a `store_name` column. Every
8//! named store shares that table; a handle binds one store name, validated
9//! once at [`Connection::kv_store`](crate::Connection::kv_store).
10//!
11//! Hyper has no native KV store and no `ON CONFLICT`/`MERGE`; `set` is an
12//! `UPDATE`-then-conditional-`INSERT` upsert. See the crate `DEVELOPMENT.md`
13//! for the design rationale.
14
15use crate::error::{Error, Result};
16
17/// Fixed backing table for every named KV store.
18///
19/// The `_hyperdb_` prefix matches the crate's internal-table convention so
20/// downstream tooling can auto-hide it from schema listings.
21pub(crate) const KV_TABLE: &str = "_hyperdb_kv_store";
22
23/// Maximum length, in bytes, of a store name or key.
24pub(crate) const KV_MAX_NAME_BYTES: usize = 512;
25
26/// Human-readable description of the allowed store-name/key charset.
27///
28/// Used in validation error messages so the allowed set is stated in one
29/// place (M-DOCUMENTED-MAGIC) rather than duplicated as a string literal.
30pub(crate) const KV_CHARSET: &str = "A-Z a-z 0-9 _ . -";
31
32/// Validates a store name or key: non-empty, ASCII `[A-Za-z0-9_.-]+`, `<= 512` bytes.
33///
34/// `kind` labels the value in the error message (`"store name"` / `"key"`).
35///
36/// # Errors
37///
38/// Returns [`Error::InvalidName`] if `name` is empty, exceeds
39/// [`KV_MAX_NAME_BYTES`] bytes, or contains a byte outside the ASCII
40/// [`KV_CHARSET`] (`A-Z a-z 0-9 _ . -`).
41pub(crate) fn validate_kv_name(name: &str, kind: &str) -> Result<()> {
42    if name.is_empty() {
43        return Err(Error::invalid_name(format!("KV {kind} must not be empty")));
44    }
45    if name.len() > KV_MAX_NAME_BYTES {
46        return Err(Error::invalid_name(format!(
47            "KV {kind} exceeds {KV_MAX_NAME_BYTES}-byte limit ({} bytes)",
48            name.len()
49        )));
50    }
51    if let Some(bad) = name
52        .bytes()
53        .find(|&b| !(b.is_ascii_alphanumeric() || b == b'_' || b == b'.' || b == b'-'))
54    {
55        return Err(Error::invalid_name(format!(
56            "KV {kind} contains an invalid byte {bad:#04x}; allowed: {KV_CHARSET}"
57        )));
58    }
59    Ok(())
60}
61
62/// Builds the `CREATE TABLE IF NOT EXISTS` DDL for the KV backing table.
63///
64/// Single source of truth for the schema, shared by the sync and async
65/// constructors and both `kv_list_stores` guards. The table has **no**
66/// `PRIMARY KEY`: Hyper rejects one at create time (`0A000: Index support is
67/// disabled`, see `hyperdb-mcp/src/table_catalog.rs`), so per-`(store_name,
68/// key)` uniqueness is an application-side invariant enforced by the
69/// UPDATE-then-conditional-INSERT upsert, not an engine constraint.
70pub(crate) fn kv_create_table_sql(table_ref: &str) -> String {
71    format!(
72        "CREATE TABLE IF NOT EXISTS {table_ref} \
73         (store_name TEXT NOT NULL, key TEXT NOT NULL, value TEXT)"
74    )
75}
76
77/// Builds the escaped `"<database>"."public"` qualifier prefix for the KV table.
78///
79/// Single source of truth for the location shape used by
80/// [`Connection::kv_store_in`](crate::Connection::kv_store_in) and
81/// [`Connection::kv_list_stores_in`](crate::Connection::kv_list_stores_in) (and
82/// their async twins): the KV table always lives in the target database's
83/// `public` schema. `database` is identifier-escaped via
84/// [`escape_name`](crate::escape_name); the fixed `public` schema name is
85/// escaped identically for symmetry.
86pub(crate) fn kv_target_prefix(database: &str) -> Result<String> {
87    Ok(format!(
88        "{}.{}",
89        crate::escape_name(database)?,
90        crate::escape_name("public")?
91    ))
92}
93
94use crate::connection::Connection;
95
96/// A handle to one named key-value store, backed by `KV_TABLE`.
97///
98/// Borrows its [`Connection`] for the handle's lifetime (`'conn`), matching
99/// the crate's [`Catalog`](crate::Catalog)/[`Inserter`](crate::Inserter)
100/// borrow convention. Open one with
101/// [`Connection::kv_store`](crate::Connection::kv_store).
102///
103/// # Examples
104///
105/// ```no_run
106/// use hyperdb_api::{Connection, CreateMode, Result};
107///
108/// fn main() -> Result<()> {
109///     let conn = Connection::connect("localhost:7483", "app.hyper", CreateMode::CreateIfNotExists)?;
110///     let kv = conn.kv_store("settings")?;
111///     kv.set("theme", "dark")?;
112///     assert_eq!(kv.get("theme")?, Some("dark".to_string()));
113///     Ok(())
114/// }
115/// ```
116#[derive(Debug)]
117pub struct KvStore<'conn> {
118    connection: &'conn Connection,
119    store_name: String,
120    table_ref: String,
121}
122
123impl<'conn> KvStore<'conn> {
124    /// Opens a handle to `name`, creating `KV_TABLE` if needed.
125    fn open(connection: &'conn Connection, name: &str, table_ref: String) -> Result<Self> {
126        validate_kv_name(name, "store name")?;
127        connection.execute_command(&kv_create_table_sql(&table_ref))?;
128        Ok(KvStore {
129            connection,
130            store_name: name.to_string(),
131            table_ref,
132        })
133    }
134
135    /// Opens a handle to a store in the default location.
136    pub(crate) fn new(connection: &'conn Connection, name: &str) -> Result<Self> {
137        Self::open(connection, name, KV_TABLE.to_string())
138    }
139
140    /// Opens a handle targeting an explicit, already-escaped table-qualifier prefix.
141    ///
142    /// Crate-internal low-level constructor behind
143    /// [`Connection::kv_store_in`](crate::Connection::kv_store_in). `target` is
144    /// interpolated directly into SQL, so the **caller must supply a pre-escaped,
145    /// SQL-safe qualifier** — public callers go through `kv_store_in`, which
146    /// escapes for them (`store_name` / `key` / `value` are always bound params,
147    /// but `target` is not).
148    pub(crate) fn with_target(
149        connection: &'conn Connection,
150        name: &str,
151        target: &str,
152    ) -> Result<Self> {
153        Self::open(connection, name, format!("{target}.{KV_TABLE}"))
154    }
155
156    /// Returns this store's validated name.
157    #[must_use]
158    pub fn name(&self) -> &str {
159        &self.store_name
160    }
161
162    /// Returns the value for `key`, or `None` if the key is absent or NULL.
163    ///
164    /// # Errors
165    ///
166    /// - [`Error::InvalidName`] if `key` is invalid.
167    /// - [`Error::FeatureNotSupported`] on gRPC transport.
168    /// - [`Error::Server`] if the query fails.
169    pub fn get(&self, key: &str) -> Result<Option<String>> {
170        validate_kv_name(key, "key")?;
171        let sql = format!(
172            "SELECT value FROM {} WHERE store_name = $1 AND key = $2",
173            self.table_ref
174        );
175        // Bind store_name/key as `&str` params (never interpolated) — uniform
176        // `&str` element types coerce cleanly to `&[&dyn ToSqlParam]`.
177        let row = self
178            .connection
179            .query_params(&sql, &[&self.store_name.as_str(), &key])?
180            .first_row()?;
181        Ok(row.and_then(|r| r.get::<String>(0)))
182    }
183
184    /// Sets `key` to `value`, inserting or overwriting (upsert).
185    ///
186    /// # Errors
187    ///
188    /// - [`Error::InvalidName`] if `key` is invalid.
189    /// - [`Error::FeatureNotSupported`] on gRPC transport.
190    /// - [`Error::Server`] if the `UPDATE`/`INSERT` fails.
191    pub fn set(&self, key: &str, value: &str) -> Result<()> {
192        validate_kv_name(key, "key")?;
193        self.upsert(key, value)
194    }
195
196    /// UPDATE-then-conditional-INSERT upsert. Assumes `key` is validated.
197    ///
198    /// Hyper has no `ON CONFLICT`; this mirrors the proven `_table_catalog`
199    /// idiom. The conditional INSERT uses distinct placeholders (`$4`/`$5`)
200    /// so it is unambiguous under the extended-query protocol.
201    fn upsert(&self, key: &str, value: &str) -> Result<()> {
202        let store = self.store_name.as_str();
203        let updated = self.connection.command_params(
204            &format!(
205                "UPDATE {} SET value = $3 WHERE store_name = $1 AND key = $2",
206                self.table_ref
207            ),
208            &[&store, &key, &value],
209        )?;
210        if updated == 0 {
211            self.connection.command_params(
212                &format!(
213                    "INSERT INTO {t} (store_name, key, value) \
214                     SELECT $1, $2, $3 \
215                     WHERE NOT EXISTS (SELECT 1 FROM {t} WHERE store_name = $4 AND key = $5)",
216                    t = self.table_ref
217                ),
218                &[&store, &key, &value, &store, &key],
219            )?;
220        }
221        Ok(())
222    }
223
224    /// Deserializes the JSON-encoded value for `key` into `T`.
225    ///
226    /// Returns `None` if the key is absent.
227    ///
228    /// # Errors
229    ///
230    /// - [`Error::InvalidName`] if `key` is invalid.
231    /// - [`Error::Serialization`] if the stored value is not valid JSON for `T`.
232    /// - [`Error::FeatureNotSupported`] / [`Error::Server`] as for [`get`](Self::get).
233    pub fn get_as<T: serde::de::DeserializeOwned>(&self, key: &str) -> Result<Option<T>> {
234        match self.get(key)? {
235            Some(json) => serde_json::from_str(&json)
236                .map(Some)
237                .map_err(|e| Error::serialization(e.to_string())),
238            None => Ok(None),
239        }
240    }
241
242    /// Serializes `value` to JSON and stores it under `key` (upsert).
243    ///
244    /// # Errors
245    ///
246    /// - [`Error::InvalidName`] if `key` is invalid.
247    /// - [`Error::Serialization`] if `value` cannot be serialized to JSON.
248    /// - [`Error::FeatureNotSupported`] / [`Error::Server`] as for [`set`](Self::set).
249    pub fn set_as<T: serde::Serialize>(&self, key: &str, value: &T) -> Result<()> {
250        validate_kv_name(key, "key")?;
251        let json = serde_json::to_string(value).map_err(|e| Error::serialization(e.to_string()))?;
252        self.upsert(key, &json)
253    }
254
255    /// Deletes `key`; returns `true` if a row was removed.
256    ///
257    /// # Errors
258    ///
259    /// - [`Error::InvalidName`] if `key` is invalid.
260    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
261    pub fn delete(&self, key: &str) -> Result<bool> {
262        validate_kv_name(key, "key")?;
263        let affected = self.connection.command_params(
264            &format!(
265                "DELETE FROM {} WHERE store_name = $1 AND key = $2",
266                self.table_ref
267            ),
268            &[&self.store_name.as_str(), &key],
269        )?;
270        Ok(affected > 0)
271    }
272
273    /// Returns whether `key` is present in this store.
274    ///
275    /// # Errors
276    ///
277    /// - [`Error::InvalidName`] if `key` is invalid.
278    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
279    pub fn exists(&self, key: &str) -> Result<bool> {
280        validate_kv_name(key, "key")?;
281        let sql = format!(
282            "SELECT 1 FROM {} WHERE store_name = $1 AND key = $2 LIMIT 1",
283            self.table_ref
284        );
285        Ok(self
286            .connection
287            .query_params(&sql, &[&self.store_name.as_str(), &key])?
288            .first_row()?
289            .is_some())
290    }
291
292    /// Returns the number of keys in this store.
293    ///
294    /// # Errors
295    ///
296    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
297    pub fn size(&self) -> Result<i64> {
298        let sql = format!(
299            "SELECT COUNT(*) FROM {} WHERE store_name = $1",
300            self.table_ref
301        );
302        // `scalar()` errors on zero rows, but COUNT(*) always returns exactly
303        // one non-NULL row, so `unwrap_or(0)` is unreachable-but-defensive.
304        Ok(self
305            .connection
306            .query_params(&sql, &[&self.store_name.as_str()])?
307            .scalar::<i64>()?
308            .unwrap_or(0))
309    }
310
311    /// Returns this store's keys, sorted ascending.
312    ///
313    /// # Errors
314    ///
315    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
316    pub fn keys(&self) -> Result<Vec<String>> {
317        let sql = format!(
318            "SELECT key FROM {} WHERE store_name = $1 ORDER BY key ASC",
319            self.table_ref
320        );
321        let mut result = self
322            .connection
323            .query_params(&sql, &[&self.store_name.as_str()])?;
324        let mut keys = Vec::new();
325        while let Some(chunk) = result.next_chunk()? {
326            for row in &chunk {
327                if let Some(k) = row.get::<String>(0) {
328                    keys.push(k);
329                }
330            }
331        }
332        Ok(keys)
333    }
334
335    /// Deletes every key in this store; returns the number removed.
336    ///
337    /// The shared backing table survives; only this store's rows are removed.
338    ///
339    /// # Errors
340    ///
341    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
342    pub fn clear(&self) -> Result<u64> {
343        self.connection.command_params(
344            &format!("DELETE FROM {} WHERE store_name = $1", self.table_ref),
345            &[&self.store_name.as_str()],
346        )
347    }
348
349    /// Removes and returns the lowest-ordered key/value pair, or `None` if empty.
350    ///
351    /// The peek and delete run in one transaction, so they apply atomically —
352    /// either both the read and the delete commit, or neither does (on error
353    /// the transaction is rolled back). A SQL-NULL value is returned as an
354    /// empty string.
355    ///
356    /// # Errors
357    ///
358    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
359    pub fn pop(&self) -> Result<Option<(String, String)>> {
360        self.connection.begin_transaction_raw()?;
361        let result = self.pop_inner();
362        match &result {
363            Ok(_) => self.connection.commit_raw()?,
364            Err(_) => {
365                // Best-effort rollback; preserve the original error.
366                let _ = self.connection.rollback_raw();
367            }
368        }
369        result
370    }
371
372    /// Transaction body for [`pop`](Self::pop).
373    fn pop_inner(&self) -> Result<Option<(String, String)>> {
374        let store = self.store_name.as_str();
375        let select = format!(
376            "SELECT key, value FROM {} WHERE store_name = $1 ORDER BY key ASC LIMIT 1",
377            self.table_ref
378        );
379        // `first_row()` consumes the `Rowset`, dropping it (and releasing its
380        // statement guard on the shared connection) BEFORE the DELETE runs —
381        // the two statements never overlap on the connection.
382        let Some(row) = self
383            .connection
384            .query_params(&select, &[&store])?
385            .first_row()?
386        else {
387            return Ok(None);
388        };
389        let key: String = row
390            .get::<String>(0)
391            .ok_or_else(|| Error::internal("kv pop: key column was unexpectedly NULL"))?;
392        let value: String = row.get::<String>(1).unwrap_or_default();
393        self.connection.command_params(
394            &format!(
395                "DELETE FROM {} WHERE store_name = $1 AND key = $2",
396                self.table_ref
397            ),
398            &[&store, &key.as_str()],
399        )?;
400        Ok(Some((key, value)))
401    }
402
403    /// Upserts every `(key, value)` pair in one transaction.
404    ///
405    /// All keys are validated before the transaction opens, so an invalid key
406    /// aborts the whole batch without writing anything.
407    ///
408    /// # Errors
409    ///
410    /// - [`Error::InvalidName`] if any key is invalid (checked before writing).
411    /// - [`Error::FeatureNotSupported`] / [`Error::Server`].
412    pub fn set_batch(&self, entries: &[(&str, &str)]) -> Result<()> {
413        for (key, _) in entries {
414            validate_kv_name(key, "key")?;
415        }
416        self.connection.begin_transaction_raw()?;
417        let result = (|| {
418            for (key, value) in entries {
419                self.upsert(key, value)?;
420            }
421            Ok(())
422        })();
423        match &result {
424            Ok(()) => self.connection.commit_raw()?,
425            Err(_) => {
426                let _ = self.connection.rollback_raw();
427            }
428        }
429        result
430    }
431}
432
433impl Connection {
434    /// Opens a handle to a named key-value store, creating the table if needed.
435    ///
436    /// # Examples
437    ///
438    /// ```no_run
439    /// # use hyperdb_api::{Connection, CreateMode, Result};
440    /// # fn example(conn: &Connection) -> Result<()> {
441    /// let kv = conn.kv_store("session")?;
442    /// # Ok(())
443    /// # }
444    /// ```
445    ///
446    /// # Errors
447    ///
448    /// - [`Error::InvalidName`] if `name` is empty, too long, or has invalid characters.
449    /// - [`Error::FeatureNotSupported`] on gRPC transport.
450    /// - [`Error::Server`] if the `CREATE TABLE IF NOT EXISTS` fails.
451    pub fn kv_store(&self, name: &str) -> Result<KvStore<'_>> {
452        KvStore::new(self, name)
453    }
454
455    /// Opens a handle to a KV store in a specific database, rather than the
456    /// default (search-path) location.
457    ///
458    /// `database` is the **unescaped** name of an attached database; the store's
459    /// backing table is created (if absent) in that database's `public` schema.
460    /// The name is identifier-escaped internally, so it is safe to pass an
461    /// arbitrary attachment alias. The store name, keys, and values are always
462    /// bound parameters.
463    ///
464    /// # Examples
465    ///
466    /// ```no_run
467    /// # use hyperdb_api::{Connection, CreateMode, Result};
468    /// # fn example(conn: &Connection) -> Result<()> {
469    /// let kv = conn.kv_store_in("persistent", "settings")?;
470    /// kv.set("theme", "dark")?;
471    /// # Ok(())
472    /// # }
473    /// ```
474    ///
475    /// # Errors
476    ///
477    /// - [`Error::InvalidName`] if `database` or `name` is invalid.
478    /// - [`Error::FeatureNotSupported`] on gRPC transport.
479    /// - [`Error::Server`] if creating the backing table fails.
480    pub fn kv_store_in(&self, database: &str, name: &str) -> Result<KvStore<'_>> {
481        KvStore::with_target(self, name, &kv_target_prefix(database)?)
482    }
483
484    /// Lists the names of every KV store that currently holds at least one key.
485    ///
486    /// Creates the backing table first (via `kv_create_table_sql`) so calling
487    /// this on a fresh database returns an empty list rather than erroring on a
488    /// missing table.
489    ///
490    /// # Errors
491    ///
492    /// - [`Error::FeatureNotSupported`] on gRPC transport.
493    /// - [`Error::Server`] if the query fails.
494    pub fn kv_list_stores(&self) -> Result<Vec<String>> {
495        self.kv_list_stores_impl(KV_TABLE)
496    }
497
498    /// Lists the KV stores that hold at least one key in a specific database.
499    ///
500    /// The location-aware companion to [`kv_list_stores`](Self::kv_list_stores):
501    /// `database` is the unescaped name of an attached database (escaped
502    /// internally). Creates the backing table in that database's `public` schema
503    /// first, so an empty database returns `[]` rather than erroring.
504    ///
505    /// # Errors
506    ///
507    /// - [`Error::InvalidName`] if `database` is invalid.
508    /// - [`Error::FeatureNotSupported`] on gRPC transport.
509    /// - [`Error::Server`] if the query fails.
510    pub fn kv_list_stores_in(&self, database: &str) -> Result<Vec<String>> {
511        let table_ref = format!("{}.{KV_TABLE}", kv_target_prefix(database)?);
512        self.kv_list_stores_impl(&table_ref)
513    }
514
515    /// Shared body for [`kv_list_stores`](Self::kv_list_stores) /
516    /// [`kv_list_stores_in`](Self::kv_list_stores_in): create the backing table
517    /// at `table_ref` (so a fresh location returns `[]`), then list the distinct
518    /// store names there. Factored out so the default and location-aware paths
519    /// cannot drift.
520    fn kv_list_stores_impl(&self, table_ref: &str) -> Result<Vec<String>> {
521        self.execute_command(&kv_create_table_sql(table_ref))?;
522        let mut result = self.execute_query(&format!(
523            "SELECT DISTINCT store_name FROM {table_ref} ORDER BY store_name ASC"
524        ))?;
525        let mut names = Vec::new();
526        while let Some(chunk) = result.next_chunk()? {
527            for row in &chunk {
528                if let Some(name) = row.get::<String>(0) {
529                    names.push(name);
530                }
531            }
532        }
533        Ok(names)
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    #[test]
542    fn accepts_valid_names() {
543        for ok in [
544            "a",
545            "store_1",
546            "my.key-2",
547            "A",
548            &"z".repeat(KV_MAX_NAME_BYTES),
549        ] {
550            assert!(validate_kv_name(ok, "key").is_ok(), "should accept {ok:?}");
551        }
552    }
553
554    #[test]
555    fn rejects_empty() {
556        let err = validate_kv_name("", "store name").unwrap_err();
557        assert!(matches!(err, Error::InvalidName(_)));
558        assert!(err.to_string().contains("must not be empty"));
559    }
560
561    #[test]
562    fn rejects_too_long() {
563        let long = "a".repeat(KV_MAX_NAME_BYTES + 1);
564        let err = validate_kv_name(&long, "key").unwrap_err();
565        assert!(matches!(err, Error::InvalidName(_)));
566        assert!(err.to_string().contains("byte limit"));
567    }
568
569    #[test]
570    fn rejects_bad_charset() {
571        for bad in ["a b", "a/b", "a'b", "a\"b", "a;b", "naïve", "a\0b"] {
572            let err = validate_kv_name(bad, "key").unwrap_err();
573            assert!(
574                matches!(err, Error::InvalidName(_)),
575                "should reject {bad:?}"
576            );
577        }
578    }
579}