slave-framework 0.13.1

//! The application's [SurrealDB](https://surrealdb.com) database.
//!
//! The app connects to a SurrealDB server (self-hosted or SurrealDB Cloud)
//! over WebSocket and runs SurrealQL through the [`surrealdb`] SDK. Several
//! processes and machines can share the database. Configure it with
//! `SURREAL_URL`, `SURREAL_NS`, `SURREAL_DB`, `SURREAL_USER`, `SURREAL_PASS`
//! and `SURREAL_AUTH_LEVEL` (see [`DatabaseConfig::from_env`]).
//!
//! Schema changes are SurrealQL files in `database/migrations/`, applied in
//! order by [`migrations`].

use std::borrow::Cow;

use surrealdb::Surreal;
use surrealdb::engine::any::Any;
use surrealdb::method::Query;
use surrealdb::opt::auth::{Database, Namespace, Root};
use surrealdb::types::QueryError;

use crate::{Error, Result};

pub mod migrations;

/// A connection to the database. Cheap to clone: clones share the connection.
pub type Db = Surreal<Any>;

/// A transaction started by [`transaction!`](crate::transaction).
pub type Transaction = surrealdb::method::Transaction<Any>;

/// Runs SurrealQL: the database ([`Db`]) or a [`Transaction`]. Write helper
/// functions against `&impl Executor` to use them in and out of transactions.
pub trait Executor {
    fn query<'a>(&'a self, query: impl Into<Cow<'a, str>>) -> Query<'a, Any>;
}

impl Executor for Db {
    fn query<'a>(&'a self, query: impl Into<Cow<'a, str>>) -> Query<'a, Any> {
        Surreal::query(self, query)
    }
}

impl Executor for Transaction {
    fn query<'a>(&'a self, query: impl Into<Cow<'a, str>>) -> Query<'a, Any> {
        Transaction::query(self, query)
    }
}

/// Errors that may be a transaction conflict, so [`transaction!`](crate::transaction)
/// knows when to retry.
pub trait MaybeConflict: From<surrealdb::Error> {
    fn is_transaction_conflict(&self) -> bool;
}

impl MaybeConflict for surrealdb::Error {
    fn is_transaction_conflict(&self) -> bool {
        is_transaction_conflict(self)
    }
}

/// Ends a transaction started by [`transaction!`](crate::transaction):
/// commits it if `result` is `Ok`, cancels it otherwise.
#[doc(hidden)]
pub async fn finish<T, E: MaybeConflict>(
    tx: Transaction,
    result: std::result::Result<T, E>,
) -> std::result::Result<T, E> {
    match result {
        Ok(value) => {
            tx.commit().await?;
            Ok(value)
        }
        Err(err) => {
            let _ = tx.cancel().await;
            Err(err)
        }
    }
}

/// Waits before retry `attempt` (2, 3, ...) of a conflicting transaction:
/// exponential backoff with jitter, so the contenders spread out.
#[doc(hidden)]
pub async fn backoff(attempt: usize) {
    let nanos = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map_or(0, |d| d.subsec_nanos() as u64);
    let base = 5u64 << attempt.min(6);
    let millis = base + nanos % base;
    tokio::time::sleep(std::time::Duration::from_millis(millis)).await;
}

/// Runs a block in a transaction: committed when it evaluates to `Ok`,
/// cancelled when it evaluates to `Err` (`?` leaves the block, not the
/// function). When the transaction conflicts with a concurrent one, the
/// block runs again in a new transaction, up to
/// [`TRANSACTION_ATTEMPTS`](database::TRANSACTION_ATTEMPTS) times.
///
/// ```ignore
/// let user = slave_framework::transaction!(&db, |tx| {
///     let user: Option<User> = tx.query("CREATE ...").await?.take(0)?;
///     tx.query("CREATE ...").await?.check()?;
///     Ok(user)
/// })?;
/// ```
///
/// The error type must implement [`MaybeConflict`](database::MaybeConflict).
/// To make two transactions that read the same rows and write different ones
/// conflict (e.g. "keep at least one owner"), have both write one shared
/// record, such as the parent row.
///
/// A macro rather than a function taking an async closure: the futures of
/// async closures are not provably `Send`, which MCP tools require.
#[macro_export]
macro_rules! transaction {
    ($db:expr, |$tx:ident| $body:block) => {{
        let db: &$crate::database::Db = $db;
        let mut attempt = 1;
        loop {
            let result = match db.clone().begin().await {
                Ok($tx) => {
                    let result = async { $body }.await;
                    $crate::database::finish($tx, result).await
                }
                Err(err) => Err(err.into()),
            };
            match result {
                Err(err)
                    if $crate::database::MaybeConflict::is_transaction_conflict(&err)
                        && attempt < $crate::database::TRANSACTION_ATTEMPTS =>
                {
                    attempt += 1;
                    $crate::database::backoff(attempt).await;
                }
                result => break result,
            }
        }
    }};
}

/// How many times [`transaction!`](crate::transaction) runs a transaction
/// that keeps failing with a conflict (see [`is_transaction_conflict`])
/// before giving up.
pub const TRANSACTION_ATTEMPTS: usize = 8;

pub const URL_ENV: &str = "SURREAL_URL";
pub const NAMESPACE_ENV: &str = "SURREAL_NS";
pub const DATABASE_ENV: &str = "SURREAL_DB";
pub const USER_ENV: &str = "SURREAL_USER";
pub const PASSWORD_ENV: &str = "SURREAL_PASS";
pub const AUTH_LEVEL_ENV: &str = "SURREAL_AUTH_LEVEL";

/// The server used when `SURREAL_URL` is not set: `surreal start` on this machine.
pub const DEFAULT_URL: &str = "ws://127.0.0.1:8000";

/// Where a SurrealDB user is defined, which decides how to sign in.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AuthLevel {
    /// A root user (`surreal start --user ... --pass ...`).
    Root,
    /// A user defined on the namespace (`DEFINE USER ... ON NAMESPACE`).
    Namespace,
    /// A user defined on the database (`DEFINE USER ... ON DATABASE`), the
    /// least privileged choice for production.
    Database,
}

impl std::str::FromStr for AuthLevel {
    type Err = String;

    fn from_str(value: &str) -> std::result::Result<Self, String> {
        match value.trim().to_ascii_lowercase().as_str() {
            "root" => Ok(Self::Root),
            "namespace" | "ns" => Ok(Self::Namespace),
            "database" | "db" => Ok(Self::Database),
            _ => Err("expected `root`, `namespace` or `database`".to_owned()),
        }
    }
}

/// Where to connect and how to sign in.
#[derive(Debug, Clone)]
pub struct DatabaseConfig {
    /// `ws://host:port` or `wss://host` (SurrealDB Cloud).
    pub url: String,
    pub namespace: String,
    pub database: String,
    /// Signs in when set; a server started with `--unauthenticated` needs none.
    pub username: Option<String>,
    pub password: Option<String>,
    pub auth_level: AuthLevel,
}

impl DatabaseConfig {
    /// Reads `SURREAL_URL` (default [`DEFAULT_URL`]), `SURREAL_NS` and
    /// `SURREAL_DB` (default `main`), `SURREAL_USER` and `SURREAL_PASS`, and
    /// `SURREAL_AUTH_LEVEL` (`root`, `namespace` or `database`; default
    /// `root`).
    pub fn from_env() -> Result<Self> {
        let non_empty = |key: &str| std::env::var(key).ok().filter(|v| !v.trim().is_empty());
        Ok(Self {
            url: crate::env::var_or(URL_ENV, DEFAULT_URL),
            namespace: crate::env::var_or(NAMESPACE_ENV, "main"),
            database: crate::env::var_or(DATABASE_ENV, "main"),
            username: non_empty(USER_ENV),
            password: non_empty(PASSWORD_ENV),
            auth_level: crate::env::parse_or(AUTH_LEVEL_ENV, AuthLevel::Root)?,
        })
    }

    /// The settings as environment variables, for child processes.
    pub fn envs(&self) -> Vec<(&'static str, String)> {
        let level = match self.auth_level {
            AuthLevel::Root => "root",
            AuthLevel::Namespace => "namespace",
            AuthLevel::Database => "database",
        };
        let mut envs = vec![
            (URL_ENV, self.url.clone()),
            (NAMESPACE_ENV, self.namespace.clone()),
            (DATABASE_ENV, self.database.clone()),
            (AUTH_LEVEL_ENV, level.to_owned()),
        ];
        envs.extend(self.username.clone().map(|v| (USER_ENV, v)));
        envs.extend(self.password.clone().map(|v| (PASSWORD_ENV, v)));
        envs
    }

    /// `url ns/db`, for logs; never includes the password.
    pub fn describe(&self) -> String {
        format!("{} {}/{}", self.url, self.namespace, self.database)
    }
}

/// Connects, signs in and selects the namespace and database.
pub async fn connect(config: &DatabaseConfig) -> Result<Db> {
    let db = surrealdb::engine::any::connect(config.url.as_str()).await?;
    if let Some(username) = &config.username {
        let password = config
            .password
            .clone()
            .ok_or(Error::MissingDatabaseSetting(PASSWORD_ENV))?;
        let username = username.clone();
        match config.auth_level {
            AuthLevel::Root => {
                db.signin(Root { username, password }).await?;
            }
            AuthLevel::Namespace => {
                db.signin(Namespace {
                    namespace: config.namespace.clone(),
                    username,
                    password,
                })
                .await?;
            }
            AuthLevel::Database => {
                db.signin(Database {
                    namespace: config.namespace.clone(),
                    database: config.database.clone(),
                    username,
                    password,
                })
                .await?;
            }
        }
    }
    db.use_ns(&config.namespace)
        .use_db(&config.database)
        .await?;
    Ok(db)
}

/// Whether `err` is a transaction that lost a race with a concurrent one;
/// running it again usually succeeds (up to [`TRANSACTION_ATTEMPTS`] times).
pub fn is_transaction_conflict(err: &surrealdb::Error) -> bool {
    matches!(err.query_details(), Some(QueryError::TransactionConflict))
        || err.message().contains("can be retried")
}

/// Whether `err` is a write rejected by a `UNIQUE` index or an existing
/// record ID.
pub fn is_already_exists(err: &surrealdb::Error) -> bool {
    err.is_already_exists() || err.message().contains("already contains")
}