ilink-hub 0.2.8

iLink-compatible multiplexer hub for WeChat ClawBot — route one WeChat account to multiple AI agent backends
Documentation
//! Database persistence layer.
//! Uses sqlx with runtime driver selection via `DATABASE_URL`:
//!   sqlite:~/.ilink-hub/ilink-hub.db → SQLite (default, file created if missing)
//!   postgres://user:pass@host/db      → PostgreSQL
//!   mysql://user:pass@host/db         → MySQL

use anyhow::Result;

use sqlx::AnyPool;

mod clients;

mod context;

mod credentials;

mod messages;

mod migrations;

mod sessions;

pub use clients::{ClientRow, HUB_DEFAULT_SENTINEL};

pub use messages::{MessageRow, RecentOutboundRow, SessionStatusEntry};

pub use sessions::BackendSessionRow;

/// Backend driver. Parsed from the URL scheme prefix in `Store::connect`
/// (sqlite: / postgres: / mysql:). The migration runner needs this to
/// pick the right SQL dialect for `try_claim_migration`, the v5 `id` column
/// clause, and the `column_exists` pre-check — `ON CONFLICT ... RETURNING`,
/// `GENERATED BY DEFAULT AS IDENTITY`, and `pragma_table_info` are
/// each driver-specific.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DatabaseKind {
    Sqlite,
    Postgres,
    MySql,
}

impl DatabaseKind {
    /// Parse the driver from the URL scheme. Accepted schemes:
    /// `sqlite:` (or empty/bare path) → SQLite,
    /// `postgres:` / `postgresql:` → PostgreSQL,
    /// `mysql:` / `mariadb:` → MySQL.
    ///
    /// Any other scheme (e.g. `postgress://`, `file:`, `http:`) returns
    /// `Err` immediately so a typo surfaces at startup rather than being
    /// silently treated as SQLite.
    fn from_url(url: &str) -> Result<Self> {
        if url.starts_with("postgres:") || url.starts_with("postgresql:") {
            #[cfg(feature = "postgres")]
            return Ok(DatabaseKind::Postgres);
            #[cfg(not(feature = "postgres"))]
            anyhow::bail!(
                "PostgreSQL support requires the `postgres` feature flag to be enabled at compile \
                 time; rebuild with --features postgres"
            );
        } else if url.starts_with("mysql:") || url.starts_with("mariadb:") {
            #[cfg(feature = "mysql")]
            return Ok(DatabaseKind::MySql);
            #[cfg(not(feature = "mysql"))]
            anyhow::bail!(
                "MySQL support requires the `mysql` feature flag to be enabled at compile time; \
                 rebuild with `--features mysql`"
            );
        } else if url.starts_with("sqlite:") || url.is_empty() {
            Ok(DatabaseKind::Sqlite)
        } else {
            if (url.starts_with('/') || url.starts_with("./") || url.starts_with("~/"))
                && !url.contains("://")
            {
                anyhow::bail!(
                    "DATABASE_URL {:?} looks like a file path; use the sqlite: scheme instead, e.g. `sqlite:{}`",
                    url, url
                );
            }
            anyhow::bail!(
                "unsupported DATABASE_URL scheme in {:?}; \
                 supported schemes: sqlite:, postgres://, postgresql://, mysql://, mariadb://",
                url
            )
        }
    }
}

pub struct Store {
    /// Write pool — always `max_connections(1)` for SQLite to serialise writes.
    pool: AnyPool,
    /// Read pool — multiple connections on SQLite WAL, same as `pool` for PG/MySQL.
    rpool: AnyPool,
    kind: DatabaseKind,
    /// Master key for encrypting/decrypting sensitive credentials (like bot tokens).
    master_key: std::sync::OnceLock<std::sync::Arc<ring::aead::LessSafeKey>>,
}

impl Store {
    /// Connect to the database and run migrations.
    ///
    /// For SQLite URLs, the database file is created automatically if it does
    /// not exist yet (equivalent to `create_if_missing(true)`).
    pub async fn connect(url: &str) -> Result<Self> {
        #[cfg(test)]
        {
            static TEST_INIT_KEY: std::sync::Once = std::sync::Once::new();
            TEST_INIT_KEY.call_once(|| {
                if std::env::var("ILINK_HUB_MASTER_KEY").is_err() {
                    std::env::set_var(
                        "ILINK_HUB_MASTER_KEY",
                        "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f",
                    );
                }
            });
        }

        sqlx::any::install_default_drivers();

        // Parse the driver from the URL scheme prefix BEFORE opening the
        // pool. The migrator needs the kind to pick driver-specific SQL
        // (INSERT...ON CONFLICT...RETURNING is SQLite/Postgres; MySQL uses
        // INSERT IGNORE; AUTOINCREMENT is SQLite-only, etc.). The URL
        // scheme is the only reliable driver signal — runtime probing via
        // `SELECT current_database()` is unreliable because both SQLite AND
        // MySQL reject that function (SQLite has no concept; MySQL uses
        // `database()` instead). See F-M3-01 in the m3 review-findings.
        let kind = DatabaseKind::from_url(url)?;
        let is_sqlite = kind == DatabaseKind::Sqlite;

        // For SQLite we must ensure the file (and its parent directory) exist
        // before connecting, because sqlx's AnyPool does not set
        // `create_if_missing` by default and will return SQLITE_CANTOPEN (14).
        if is_sqlite {
            let url_owned = url.to_string();
            tokio::time::timeout(
                std::time::Duration::from_secs(10),
                tokio::task::spawn_blocking(move || Self::ensure_sqlite_file(&url_owned)),
            )
            .await
            .map_err(|_| anyhow::anyhow!("ensure_sqlite_file timed out after 10s"))?
            .map_err(|e| anyhow::anyhow!("spawn_blocking failed: {e}"))??;
        }

        // SQLite connection pool strategy:
        //
        // Write pool (max_connections=1): serialises all writes and DDL.  A single
        // writer is required because SQLite's file-level write lock means concurrent
        // writers from separate connections will see SQLITE_BUSY; WAL mode reduces
        // but does not eliminate this for concurrent writers.
        //
        // Read pool (max_connections=4): WAL mode allows multiple concurrent readers
        // without blocking the writer — each reader snapshots the WAL at its own
        // read mark.  For :memory: databases reads and writes share the single write
        // connection (the read pool is set to the same pool) because each physical
        // connection to ":memory:" creates an independent empty database.
        //
        // Non-SQLite (Postgres, MySQL): both pools point to the same AnyPool;
        // those engines handle their own concurrency.
        //
        // We set `journal_mode=WAL` and `busy_timeout=5000` on every connection
        // so behaviour is explicit rather than relying on sqlx defaults.
        // For :memory: SQLite, reads and writes share the single write connection —
        // separate physical connections each see an independent empty database.
        let is_memory_sqlite = is_sqlite && (url.contains(":memory:") || url.ends_with("sqlite:"));

        let (pool, rpool) = if is_sqlite {
            let wpool = sqlx::pool::PoolOptions::<sqlx::Any>::new()
                .max_connections(1)
                .after_connect(|conn, _meta| {
                    Box::pin(async move {
                        // busy_timeout must be set first so that the subsequent
                        // journal_mode=WAL switch (which requires an exclusive
                        // file lock) will wait up to 5s instead of returning
                        // SQLITE_BUSY immediately.
                        sqlx::query("PRAGMA busy_timeout=5000; PRAGMA journal_mode=WAL")
                            .execute(&mut *conn)
                            .await?;
                        Ok(())
                    })
                })
                .connect(url)
                .await?;
            let rpool = if is_memory_sqlite {
                wpool.clone()
            } else {
                sqlx::pool::PoolOptions::<sqlx::Any>::new()
                    .max_connections(4)
                    .after_connect(|conn, _meta| {
                        Box::pin(async move {
                            sqlx::query("PRAGMA busy_timeout=5000; PRAGMA journal_mode=WAL")
                                .execute(&mut *conn)
                                .await?;
                            Ok(())
                        })
                    })
                    .connect(url)
                    .await?
            };
            (wpool, rpool)
        } else {
            let pool = AnyPool::connect(url).await?;
            (pool.clone(), pool)
        };

        let store = Self {
            pool,
            rpool,
            kind,
            master_key: std::sync::OnceLock::new(),
        };
        store.run_migrations().await?;
        Ok(store)
    }

    pub fn set_master_key(
        &self,
        key: std::sync::Arc<ring::aead::LessSafeKey>,
    ) -> Result<(), std::sync::Arc<ring::aead::LessSafeKey>> {
        self.master_key.set(key)
    }

    pub fn master_key(&self) -> Option<&std::sync::Arc<ring::aead::LessSafeKey>> {
        self.master_key.get()
    }

    #[cfg(test)]
    pub fn pool(&self) -> &sqlx::AnyPool {
        &self.pool
    }

    /// Extract the file path from a SQLite URL and create the file + parent
    /// directories if they do not already exist.
    fn ensure_sqlite_file(url: &str) -> Result<()> {
        // Strip the "sqlite:" scheme prefix; handle the optional // or ///
        let path_part = url
            .strip_prefix("sqlite:///")
            .or_else(|| url.strip_prefix("sqlite://"))
            .or_else(|| url.strip_prefix("sqlite:"))
            .unwrap_or("");

        // Drop any query string (e.g. "?mode=rwc")
        let path_str = path_part.split('?').next().unwrap_or("").trim();

        // Skip in-memory databases (:memory: or empty)
        if path_str.is_empty() || path_str == ":memory:" {
            return Ok(());
        }

        let path = std::path::Path::new(path_str);
        if let Some(parent) = path.parent() {
            if !parent.as_os_str().is_empty() {
                std::fs::create_dir_all(parent)?;
            }
        }
        if !path.exists() {
            match std::fs::OpenOptions::new()
                .create_new(true)
                .write(true)
                .open(path)
            {
                Ok(_) => {}
                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
                    // Another process created the file between our `exists()`
                    // check and `create_new()`; the file exists now — no
                    // truncation, no data loss.
                }
                Err(e) => return Err(e.into()),
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod store_tests;