plenum-bot/src/key_value.rs

use rusqlite::{Connection, Result, params};
use std::ops::Index;
use std::cell::Cell;
use std::collections::HashSet;

/// Simple SQLite-backed key/value store.
///
/// All failing writes will log a message, any errors will set `has_errors`.
///
/// There's `Index` support for convenience (`&kvstore["foo"]`),
/// but **it will leak data.**  
/// (That's perfectly fine when used for getting data once for the duration of
/// the program, just use `.get()` for stuff that changes.)
pub struct KeyValueStore {
    conn: Connection,
    has_write_errors: Cell<bool>,
}

impl KeyValueStore {
    /// Open or create a DB with a `kv_store` table storing all tuples.
    pub fn new(db_path: &str) -> Result<Self> {
        let conn = Connection::open(db_path)?;
        conn.execute(
            "CREATE TABLE IF NOT EXISTS kv_store (
                  key TEXT PRIMARY KEY,
                  value TEXT NOT NULL
              )",
            [],
        )?;
        Ok(Self { conn, has_write_errors: Cell::new(false), })
    }

    /// Report if any write errors occurred.
    pub fn has_errors(&self) -> bool {
        self.has_write_errors.get()
    }

    /// Get value at `key`.  (Use `.unwrap_or(default)` & friends to work with the Result.)
    pub fn get(&self, key: &str) -> Result<String> {
        let mut stmt = self.conn.prepare("SELECT value FROM kv_store WHERE key = ?1")?;
        let mut rows = stmt.query(params![key])?;

        if let Some(row) = rows.next()? {
            let result = row.get(0)?;
            Ok(result)
        } else {
            Err(rusqlite::Error::QueryReturnedNoRows)
        }
    }

    /// If `key` is absent from the DB, write `key`/`value` pair (else leave unchanged).
    ///
    /// Will panic on error. (The idea is to call this a bunch of times at boot,
    /// to populate the DB if missing / newly created, and not use it after.)
    pub fn default(&self, key: &str, value: &str) {
        self.conn.execute(
            "INSERT INTO kv_store (key, value) VALUES (?1, ?2)
                 ON CONFLICT(key) DO NOTHING",
            params![key, value],
        ).expect(&format!("Failed to write default at key: {}", key));
    }

    /// Write a `key`/`value` pair to the DB.
    ///
    /// Errors log the failed write and set `has_errors`.
    pub fn set(&self, key: &str, value: &str) -> Result<()> {
        match self.conn.execute(
            "INSERT INTO kv_store (key, value) VALUES (?1, ?2)
             ON CONFLICT(key) DO UPDATE SET value = excluded.value",
            params![key, value],
        ) {
            Ok(_) => Ok(()),
            Err(e) => {
                eprintln!("Failed DB write: ({}, {})", key, value);
                self.has_write_errors.set(true);
                Err(e)
            }
        }
    }

    /// Delete the key-value pair associated with `key`.
    ///
    /// Errors log the failed write and set `has_errors`.
    pub fn delete(&self, key: &str) -> Result<()> {
        match self.conn.execute("DELETE FROM kv_store WHERE key = ?1", params![key]) {
            Ok(_) => Ok(()),
            Err(e) => {
                eprintln!("Failed DB write: ({}, NULL)", key);
                self.has_write_errors.set(true);
                Err(e)
            }
        }
    }

    /// Dump the current DB contents to stderr.
    ///
    /// For all keys listed in `redacted_keys`, no value is printed, only the presence is reported.
    pub fn dump_redacting(&self, redacted_keys: &[&str]) -> Result<()> {
        let exclude_set: HashSet<&str> = redacted_keys.iter().cloned().collect();
        let mut stmt = self.conn.prepare("SELECT key, value FROM kv_store")?;
        let mut rows = stmt.query(params![])?;
        while let Some(row) = rows.next()? {
            let key: String = row.get(0)?;
            if exclude_set.contains(key.as_str()) {
                eprintln!("{} = REDACTED", key);
            } else {
                let value: String = row.get(1)?;
                eprintln!("{} = {}", key, value);
            }
        }
        Ok(())
    }
}

impl Index<&str> for KeyValueStore {
    type Output = str;

    fn index(&self, key: &str) -> &Self::Output {
        Box::leak(Box::new(self.get(key).unwrap())).as_str()
    }
}