foukoapi 0.1.2-alpha.2

//! Pluggable key-value storage.
//!
//! A [`Storage`] is any async key-value store. FoukoApi ships two
//! implementations out of the box:
//!
//! - [`MemoryStorage`] - non-persistent, perfect for tests and examples.
//! - [`SqliteStorage`] (feature `sqlite`, on by default) - a tiny bundled
//!   SQLite file, auto-created if missing. Great default for small bots.
//!
//! External backends (Postgres, Redis, etc.) are out of scope for the core
//! crate; implement [`Storage`] yourself and pass it to
//! [`crate::Accounts::new`] / your handlers.

use crate::{config::DbUrl, Error, Result};
use async_trait::async_trait;
use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};

/// Key-value storage used by FoukoApi (account linking, per-user state, ...).
#[async_trait]
pub trait Storage: Send + Sync + 'static {
    /// Fetch the value stored under `key`, or `None` if absent.
    async fn get(&self, key: &str) -> Result<Option<String>>;
    /// Write `value` under `key`, overwriting any previous value.
    async fn set(&self, key: &str, value: &str) -> Result<()>;
    /// Remove a key. It is not an error if the key didn't exist.
    async fn del(&self, key: &str) -> Result<()>;
    /// Write `value` under `key` only if the key is absent. Returns `true`
    /// when the write happened, `false` when the key already existed.
    ///
    /// The default implementation is a get-then-set and is NOT atomic:
    /// two concurrent callers can both observe the key missing and both
    /// "win". Backends should override it with a real atomic upsert (the
    /// bundled memory and SQLite stores do).
    async fn set_nx(&self, key: &str, value: &str) -> Result<bool> {
        if self.get(key).await?.is_some() {
            return Ok(false);
        }
        self.set(key, value).await?;
        Ok(true)
    }
    /// Every `(key, value)` pair whose key starts with `prefix`.
    ///
    /// Order is unspecified. The default implementation returns nothing,
    /// so custom backends keep compiling; the bundled memory and SQLite
    /// stores override it with a real scan. Handy for iterating a set of
    /// related keys (open polls, pending reminders, a leaderboard index)
    /// without keeping a separate list by hand.
    async fn list_prefix(&self, _prefix: &str) -> Result<Vec<(String, String)>> {
        Ok(Vec::new())
    }
}

/// Ready-to-use storage handle, erased behind an `Arc`.
///
/// Returned by [`open_storage`] so a bot's `main.rs` can stay short.
pub type AnyStorage = Arc<dyn Storage>;

/// Open whatever storage `FOUKO_DB` points at.
///
/// - `sqlite:/path.db` (or empty/`memory:`): local file created if missing /
///   in-memory.
/// - `postgres://...` and other URL schemes are returned as
///   [`Error::Other`] for now - implement your own [`Storage`] impl.
pub fn open_storage() -> Result<AnyStorage> {
    open_storage_from(DbUrl::from_env()?)
}

/// Like [`open_storage`] but you get to pass the URL directly.
pub fn open_storage_from(url: DbUrl) -> Result<AnyStorage> {
    match url {
        DbUrl::Memory => Ok(Arc::new(MemoryStorage::new())),
        #[cfg(feature = "sqlite")]
        DbUrl::Sqlite(path) => Ok(Arc::new(SqliteStorage::open(&path)?)),
        #[cfg(not(feature = "sqlite"))]
        DbUrl::Sqlite(_) => Err(Error::Other(
            "sqlite backend is not enabled - build with the `sqlite` feature".into(),
        )),
        DbUrl::External(url) => Err(Error::Other(format!(
            "FoukoApi does not bundle a driver for {url}. Implement Storage yourself and plug it in."
        ))),
    }
}

// ---------- Memory backend ---------------------------------------------------

/// In-memory, non-persistent storage. Lost on restart.
///
/// Useful in tests, examples, and for bots that genuinely don't need state.
#[derive(Debug, Clone, Default)]
pub struct MemoryStorage {
    inner: Arc<Mutex<HashMap<String, String>>>,
}

impl MemoryStorage {
    /// New empty store.
    pub fn new() -> Self {
        Self::default()
    }

    /// Lock the map, recovering from a poisoned mutex. The data is just
    /// strings, so a panic mid-operation can't leave it inconsistent.
    fn map(&self) -> std::sync::MutexGuard<'_, HashMap<String, String>> {
        self.inner.lock().unwrap_or_else(|p| p.into_inner())
    }
}

#[async_trait]
impl Storage for MemoryStorage {
    async fn get(&self, key: &str) -> Result<Option<String>> {
        Ok(self.map().get(key).cloned())
    }
    async fn set(&self, key: &str, value: &str) -> Result<()> {
        self.map().insert(key.to_owned(), value.to_owned());
        Ok(())
    }
    async fn del(&self, key: &str) -> Result<()> {
        self.map().remove(key);
        Ok(())
    }
    async fn set_nx(&self, key: &str, value: &str) -> Result<bool> {
        // Atomic: entry() decides and inserts under one lock.
        let mut map = self.map();
        match map.entry(key.to_owned()) {
            std::collections::hash_map::Entry::Occupied(_) => Ok(false),
            std::collections::hash_map::Entry::Vacant(e) => {
                e.insert(value.to_owned());
                Ok(true)
            }
        }
    }
    async fn list_prefix(&self, prefix: &str) -> Result<Vec<(String, String)>> {
        Ok(self
            .map()
            .iter()
            .filter(|(k, _)| k.starts_with(prefix))
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect())
    }
}

// ---------- SQLite backend ---------------------------------------------------

#[cfg(feature = "sqlite")]
#[cfg_attr(docsrs, doc(cfg(feature = "sqlite")))]
mod sqlite_impl {
    use super::*;
    use rusqlite::{params, Connection};
    use std::path::Path;
    use std::sync::Mutex as StdMutex;

    /// Persistent SQLite-backed storage.
    ///
    /// `SqliteStorage::open(path)` auto-creates the file and table if they
    /// don't exist, so a first-run bot just works.
    pub struct SqliteStorage {
        conn: StdMutex<Connection>,
    }

    impl SqliteStorage {
        /// Open (or create) a SQLite file at `path`.
        pub fn open(path: &Path) -> Result<Self> {
            if let Some(parent) = path.parent() {
                if !parent.as_os_str().is_empty() {
                    std::fs::create_dir_all(parent).map_err(|e| {
                        Error::Other(format!("creating {} dir: {e}", parent.display()))
                    })?;
                }
            }
            let conn = Connection::open(path)
                .map_err(|e| Error::Other(format!("opening sqlite {}: {e}", path.display())))?;
            conn.execute_batch(
                "CREATE TABLE IF NOT EXISTS foukoapi_kv (
                    k TEXT PRIMARY KEY,
                    v TEXT NOT NULL
                 );",
            )
            .map_err(|e| Error::Other(format!("creating kv table: {e}")))?;
            Ok(Self {
                conn: StdMutex::new(conn),
            })
        }
    }

    #[async_trait]
    impl Storage for SqliteStorage {
        async fn get(&self, key: &str) -> Result<Option<String>> {
            let conn = self
                .conn
                .lock()
                .map_err(|_| Error::Other("sqlite mutex poisoned".into()))?;
            Ok(conn
                .query_row(
                    "SELECT v FROM foukoapi_kv WHERE k = ?1",
                    params![key],
                    |row| row.get::<_, String>(0),
                )
                .ok())
        }
        async fn set(&self, key: &str, value: &str) -> Result<()> {
            let conn = self
                .conn
                .lock()
                .map_err(|_| Error::Other("sqlite mutex poisoned".into()))?;
            conn.execute(
                "INSERT INTO foukoapi_kv (k, v) VALUES (?1, ?2)
                 ON CONFLICT(k) DO UPDATE SET v = excluded.v",
                params![key, value],
            )
            .map_err(|e| Error::Other(format!("sqlite set: {e}")))?;
            Ok(())
        }
        async fn del(&self, key: &str) -> Result<()> {
            let conn = self
                .conn
                .lock()
                .map_err(|_| Error::Other("sqlite mutex poisoned".into()))?;
            conn.execute("DELETE FROM foukoapi_kv WHERE k = ?1", params![key])
                .map_err(|e| Error::Other(format!("sqlite del: {e}")))?;
            Ok(())
        }
        async fn set_nx(&self, key: &str, value: &str) -> Result<bool> {
            let conn = self
                .conn
                .lock()
                .map_err(|_| Error::Other("sqlite mutex poisoned".into()))?;
            // Atomic: the conflict clause makes insert-if-absent a single
            // statement; changes() tells us whether the row went in.
            let changed = conn
                .execute(
                    "INSERT INTO foukoapi_kv (k, v) VALUES (?1, ?2)
                     ON CONFLICT(k) DO NOTHING",
                    params![key, value],
                )
                .map_err(|e| Error::Other(format!("sqlite set_nx: {e}")))?;
            Ok(changed > 0)
        }
        async fn list_prefix(&self, prefix: &str) -> Result<Vec<(String, String)>> {
            let conn = self
                .conn
                .lock()
                .map_err(|_| Error::Other("sqlite mutex poisoned".into()))?;
            // Escape LIKE wildcards so a prefix containing `%` or `_` is
            // matched literally, then append `%` for the actual prefix
            // match.
            let pattern = format!(
                "{}%",
                prefix
                    .replace('\\', "\\\\")
                    .replace('%', "\\%")
                    .replace('_', "\\_")
            );
            let mut stmt = conn
                .prepare("SELECT k, v FROM foukoapi_kv WHERE k LIKE ?1 ESCAPE '\\'")
                .map_err(|e| Error::Other(format!("sqlite list_prefix prepare: {e}")))?;
            let rows = stmt
                .query_map(params![pattern], |row| {
                    Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
                })
                .map_err(|e| Error::Other(format!("sqlite list_prefix query: {e}")))?;
            let mut out = Vec::new();
            for row in rows {
                out.push(row.map_err(|e| Error::Other(format!("sqlite list_prefix row: {e}")))?);
            }
            Ok(out)
        }
    }
}

#[cfg(feature = "sqlite")]
pub use sqlite_impl::SqliteStorage;

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn memory_list_prefix() {
        let s = MemoryStorage::new();
        s.set("a:1", "x").await.unwrap();
        s.set("a:2", "y").await.unwrap();
        s.set("b:1", "z").await.unwrap();
        let mut found = s.list_prefix("a:").await.unwrap();
        found.sort();
        assert_eq!(
            found,
            vec![
                ("a:1".to_owned(), "x".to_owned()),
                ("a:2".to_owned(), "y".to_owned())
            ]
        );
        assert!(s.list_prefix("nope:").await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn memory_set_nx() {
        let s = MemoryStorage::new();
        assert!(s.set_nx("k", "first").await.unwrap());
        assert!(!s.set_nx("k", "second").await.unwrap());
        assert_eq!(s.get("k").await.unwrap().as_deref(), Some("first"));
    }

    #[cfg(feature = "sqlite")]
    #[tokio::test]
    async fn sqlite_set_nx() {
        let dir = std::env::temp_dir().join(format!("foukoapi-setnx-{}", std::process::id()));
        let path = dir.join("kv.db");
        let s = SqliteStorage::open(&path).unwrap();
        assert!(s.set_nx("k", "first").await.unwrap());
        assert!(!s.set_nx("k", "second").await.unwrap());
        assert_eq!(s.get("k").await.unwrap().as_deref(), Some("first"));
        drop(s);
        let _ = std::fs::remove_dir_all(&dir);
    }
}