orch/store/
kv.rs

1use super::*;
2
3impl TaskStore {
4    pub async fn kv_get(&self, key: &str) -> anyhow::Result<Option<String>> {
5        let result: Option<(String,)> = sqlx::query_as("SELECT value FROM kv WHERE key = ?")
6            .bind(key)
7            .fetch_optional(&self.pool)
8            .await?;
9        Ok(result.map(|r| r.0))
10    }
11
12    /// Set a value in the KV store (upsert).
13    pub async fn kv_set(&self, key: &str, value: &str) -> anyhow::Result<()> {
14        sqlx::query(
15        "INSERT INTO kv (key, value, updated_at) VALUES (?, ?, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'))
16         ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at",
17    )
18    .bind(key)
19    .bind(value)
20    .execute(&self.pool)
21    .await?;
22        Ok(())
23    }
24
25    /// Atomically increment an integer counter in the KV store, inserting it as 1 if absent.
26    ///
27    /// Executes a single SQL statement (`INSERT … ON CONFLICT … DO UPDATE`) so concurrent
28    /// callers cannot race — SQLite's serialised write lock ensures the read-modify-write is
29    /// indivisible.  Returns the new value after the increment.
30    pub async fn kv_increment(&self, key: &str) -> anyhow::Result<u64> {
31        let row: (i64,) = sqlx::query_as(
32            "INSERT INTO kv (key, value, updated_at) VALUES (?1, '1', strftime('%Y-%m-%dT%H:%M:%SZ', 'now'))
33             ON CONFLICT(key) DO UPDATE
34               SET value = CAST(value AS INTEGER) + 1,
35                   updated_at = strftime('%Y-%m-%dT%H:%M:%SZ', 'now')
36             RETURNING CAST(value AS INTEGER)",
37        )
38        .bind(key)
39        .fetch_one(&self.pool)
40        .await?;
41        // The SQL always returns a value >= 1 (insert path returns 1, update adds 1 to
42        // the previous integer value). Previously this was clamped with `.max(1)`,
43        // which silently masked database corruption or unexpected negative values.
44        // Use a debug assertion to loudly detect unexpected values in debug builds
45        // while returning the actual stored value in release builds.
46        debug_assert!(
47            row.0 >= 1,
48            "kv_increment returned unexpected value {}",
49            row.0
50        );
51        Ok(row.0 as u64)
52    }
53
54    /// Insert a value only if the key is absent, then return the stored winner.
55    ///
56    /// Uses `INSERT OR IGNORE` so concurrent callers racing to create the same key
57    /// both converge on the same stored value — whoever lands first wins, and the
58    /// loser's value is silently discarded.  The unconditional `SELECT` afterwards
59    /// always reads back whichever value is actually in the store.
60    pub async fn kv_insert_if_absent(&self, key: &str, value: &str) -> anyhow::Result<String> {
61        sqlx::query(
62            "INSERT OR IGNORE INTO kv (key, value, updated_at)
63             VALUES (?, ?, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'))",
64        )
65        .bind(key)
66        .bind(value)
67        .execute(&self.pool)
68        .await?;
69        let (stored,): (String,) = sqlx::query_as("SELECT value FROM kv WHERE key = ?")
70            .bind(key)
71            .fetch_one(&self.pool)
72            .await?;
73        Ok(stored)
74    }
75
76    /// Delete a key from the KV store.
77    pub async fn kv_delete(&self, key: &str) -> anyhow::Result<()> {
78        sqlx::query("DELETE FROM kv WHERE key = ?")
79            .bind(key)
80            .execute(&self.pool)
81            .await?;
82        Ok(())
83    }
84
85    /// List all (key, value) pairs where the key starts with `prefix`.
86    ///
87    /// The prefix is escaped so that literal `%` and `_` characters are treated as
88    /// plain text rather than SQLite LIKE wildcards.
89    pub async fn kv_list_prefix(&self, prefix: &str) -> anyhow::Result<Vec<(String, String)>> {
90        // Escape LIKE metacharacters so a prefix containing '%' or '_' (or the chosen
91        // escape char '\') is matched literally.
92        let escaped = prefix
93            .replace('\\', "\\\\")
94            .replace('%', "\\%")
95            .replace('_', "\\_");
96        let pattern = format!("{escaped}%");
97        let rows: Vec<(String, String)> =
98            sqlx::query_as("SELECT key, value FROM kv WHERE key LIKE ? ESCAPE '\\'")
99                .bind(&pattern)
100                .fetch_all(&self.pool)
101                .await?;
102        Ok(rows)
103    }
104}
orch/store/kv.rs

orch/store/
kv.rs
