secrets_storage_postgres/
lib.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use secrets_core::storage::{StorageBackend, StorageEntry, StorageError, StorageResult};
5use sha2::{Digest, Sha256};
6use sqlx::pool::PoolConnection;
7use sqlx::{PgPool, Postgres};
8use tokio::sync::Mutex;
9
10pub struct PgStorage {
11    pool: PgPool,
12    /// Connections held open purely to keep `pg_advisory_lock` sessions alive.
13    ///
14    /// A Postgres advisory lock belongs to the *session* that took it, so the
15    /// connection cannot go back to the pool: another query reusing it could
16    /// release the lock, and a pooled `pg_advisory_unlock` might run on a
17    /// connection that never held it. Holding the connection also gives
18    /// failover for free — when the process dies the connection closes and
19    /// Postgres drops the lock, with no lease timeout to tune.
20    locks: Mutex<HashMap<String, PoolConnection<Postgres>>>,
21}
22
23impl PgStorage {
24    /// Connect and apply this crate's migrations (`src/migrations/`), the
25    /// behaviour every release before 1.1.0 had. Equivalent to
26    /// `connect_with(database_url, true)`.
27    pub async fn connect(database_url: &str) -> Result<Self, sqlx::Error> {
28        Self::connect_with(database_url, true).await
29    }
30
31    /// Connect, applying the migrations only if `migrate` is true.
32    ///
33    /// `migrate = false` is for deployments where something else owns the
34    /// schema — `typednotes-infra` applies `src/migrations/` as a declared
35    /// history, so the SQL is reviewed in a plan before it runs and the
36    /// serving identity needs no DDL rights. The schema is then *checked*
37    /// rather than assumed: a missing `kv_store` fails here, at startup,
38    /// naming the cause, instead of as a confusing error on the first
39    /// request. Note that the two modes do not mix on one database: sqlx
40    /// records what it applied in `_sqlx_migrations`, an external runner does
41    /// not, so switching an existing database from one to the other means
42    /// reconciling that bookkeeping first.
43    pub async fn connect_with(database_url: &str, migrate: bool) -> Result<Self, sqlx::Error> {
44        let pool = PgPool::connect(database_url).await?;
45        if migrate {
46            sqlx::migrate!("./src/migrations").run(&pool).await?;
47        } else {
48            sqlx::query("SELECT 1 FROM kv_store LIMIT 0")
49                .execute(&pool)
50                .await
51                .map_err(|e| {
52                    sqlx::Error::Configuration(
53                        format!(
54                            "storage migrations are disabled, but the schema is not in \
55                             place ({e}); apply src/migrations/ externally first, or \
56                             enable migrations"
57                        )
58                        .into(),
59                    )
60                })?;
61        }
62        Ok(Self {
63            pool,
64            locks: Mutex::new(HashMap::new()),
65        })
66    }
67}
68
69/// Advisory locks are keyed by a single 64-bit integer, so a name has to be
70/// folded down to one. It must be stable across processes and releases —
71/// `DefaultHasher` is explicitly not, so hash with SHA-256 instead.
72fn advisory_lock_id(key: &str) -> i64 {
73    let digest = Sha256::digest(key.as_bytes());
74    i64::from_be_bytes(digest[..8].try_into().expect("sha256 digest is 32 bytes"))
75}
76
77#[async_trait]
78impl StorageBackend for PgStorage {
79    async fn get(&self, path: &str) -> StorageResult<Option<StorageEntry>> {
80        let row = sqlx::query_as::<_, (Vec<u8>, Option<chrono::DateTime<chrono::Utc>>)>(
81            "SELECT value, expires_at FROM kv_store WHERE path = $1",
82        )
83        .bind(path)
84        .fetch_optional(&self.pool)
85        .await
86        .map_err(|e| StorageError::Backend(e.to_string()))?;
87
88        Ok(row.map(|(value, expires_at)| StorageEntry { value, expires_at }))
89    }
90
91    async fn put(&self, path: &str, entry: StorageEntry) -> StorageResult<()> {
92        sqlx::query(
93            "INSERT INTO kv_store (path, value, expires_at, updated_at)
94             VALUES ($1, $2, $3, now())
95             ON CONFLICT (path) DO UPDATE SET value = $2, expires_at = $3, updated_at = now()",
96        )
97        .bind(path)
98        .bind(entry.value)
99        .bind(entry.expires_at)
100        .execute(&self.pool)
101        .await
102        .map_err(|e| StorageError::Backend(e.to_string()))?;
103        Ok(())
104    }
105
106    async fn delete(&self, path: &str) -> StorageResult<()> {
107        sqlx::query("DELETE FROM kv_store WHERE path = $1")
108            .bind(path)
109            .execute(&self.pool)
110            .await
111            .map_err(|e| StorageError::Backend(e.to_string()))?;
112        Ok(())
113    }
114
115    async fn ping(&self) -> StorageResult<()> {
116        sqlx::query("SELECT 1")
117            .execute(&self.pool)
118            .await
119            .map_err(|e| StorageError::Backend(e.to_string()))?;
120        Ok(())
121    }
122
123    /// Served by the partial index on `expires_at` from the initial
124    /// migration, so the reaper no longer reads every lease to find the few
125    /// that expired.
126    async fn list_expired(
127        &self,
128        prefix: &str,
129        now: chrono::DateTime<chrono::Utc>,
130    ) -> StorageResult<Vec<String>> {
131        let rows: Vec<(String,)> = sqlx::query_as(
132            "SELECT path FROM kv_store
133             WHERE path LIKE $1 || '%' AND expires_at IS NOT NULL AND expires_at <= $2",
134        )
135        .bind(prefix)
136        .bind(now)
137        .fetch_all(&self.pool)
138        .await
139        .map_err(|e| StorageError::Backend(e.to_string()))?;
140        Ok(rows.into_iter().map(|(p,)| p).collect())
141    }
142
143    /// One statement, so the comparison and the write cannot be interleaved.
144    async fn replace_if_unchanged(
145        &self,
146        path: &str,
147        expected: &[u8],
148        entry: StorageEntry,
149    ) -> StorageResult<bool> {
150        let result = sqlx::query(
151            "UPDATE kv_store SET value = $3, expires_at = $4, updated_at = now()
152             WHERE path = $1 AND value = $2",
153        )
154        .bind(path)
155        .bind(expected)
156        .bind(entry.value)
157        .bind(entry.expires_at)
158        .execute(&self.pool)
159        .await
160        .map_err(|e| StorageError::Backend(e.to_string()))?;
161        Ok(result.rows_affected() == 1)
162    }
163
164    async fn try_acquire_lock(&self, key: &str) -> StorageResult<bool> {
165        let mut locks = self.locks.lock().await;
166        // Already leading. Re-taking the same advisory lock on the same
167        // session would succeed and just bump Postgres' own counter, so skip
168        // the round trip entirely.
169        if locks.contains_key(key) {
170            return Ok(true);
171        }
172
173        let mut conn = self
174            .pool
175            .acquire()
176            .await
177            .map_err(|e| StorageError::Backend(e.to_string()))?;
178        let (acquired,): (bool,) = sqlx::query_as("SELECT pg_try_advisory_lock($1)")
179            .bind(advisory_lock_id(key))
180            .fetch_one(&mut *conn)
181            .await
182            .map_err(|e| StorageError::Backend(e.to_string()))?;
183
184        if acquired {
185            locks.insert(key.to_string(), conn);
186        }
187        Ok(acquired)
188    }
189
190    async fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
191        let rows: Vec<(String,)> =
192            sqlx::query_as("SELECT path FROM kv_store WHERE path LIKE $1 || '%'")
193                .bind(prefix)
194                .fetch_all(&self.pool)
195                .await
196                .map_err(|e| StorageError::Backend(e.to_string()))?;
197        Ok(rows.into_iter().map(|(p,)| p).collect())
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use super::advisory_lock_id;
204
205    /// The id has to be identical in every process that competes for the
206    /// lock, so this is a pinned value, not a round-trip check.
207    #[test]
208    fn advisory_lock_ids_are_stable_and_distinct() {
209        assert_eq!(advisory_lock_id("secrets/lease-reaper"), advisory_lock_id("secrets/lease-reaper"));
210        assert_ne!(advisory_lock_id("secrets/lease-reaper"), advisory_lock_id("secrets/other"));
211    }
212}
secrets_storage_postgres/lib.rs

secrets_storage_postgres/
lib.rs
