slave_framework/
database.rs

1//! The application's [SurrealDB](https://surrealdb.com) database.
2//!
3//! The app connects to a SurrealDB server (self-hosted or SurrealDB Cloud)
4//! over WebSocket and runs SurrealQL through the [`surrealdb`] SDK. Several
5//! processes and machines can share the database. Configure it with
6//! `SURREAL_URL`, `SURREAL_NS`, `SURREAL_DB`, `SURREAL_USER`, `SURREAL_PASS`
7//! and `SURREAL_AUTH_LEVEL` (see [`DatabaseConfig::from_env`]).
8//!
9//! Schema changes are SurrealQL files in `database/migrations/`, applied in
10//! order by [`migrations`].
11
12use std::borrow::Cow;
13
14use surrealdb::Surreal;
15use surrealdb::engine::any::Any;
16use surrealdb::method::Query;
17use surrealdb::opt::auth::{Database, Namespace, Root};
18use surrealdb::types::QueryError;
19
20use crate::{Error, Result};
21
22pub mod migrations;
23
24/// A connection to the database. Cheap to clone: clones share the connection.
25pub type Db = Surreal<Any>;
26
27/// A transaction started by [`transaction!`](crate::transaction).
28pub type Transaction = surrealdb::method::Transaction<Any>;
29
30/// Runs SurrealQL: the database ([`Db`]) or a [`Transaction`]. Write helper
31/// functions against `&impl Executor` to use them in and out of transactions.
32pub trait Executor {
33    fn query<'a>(&'a self, query: impl Into<Cow<'a, str>>) -> Query<'a, Any>;
34}
35
36impl Executor for Db {
37    fn query<'a>(&'a self, query: impl Into<Cow<'a, str>>) -> Query<'a, Any> {
38        Surreal::query(self, query)
39    }
40}
41
42impl Executor for Transaction {
43    fn query<'a>(&'a self, query: impl Into<Cow<'a, str>>) -> Query<'a, Any> {
44        Transaction::query(self, query)
45    }
46}
47
48/// Errors that may be a transaction conflict, so [`transaction!`](crate::transaction)
49/// knows when to retry.
50pub trait MaybeConflict: From<surrealdb::Error> {
51    fn is_transaction_conflict(&self) -> bool;
52}
53
54impl MaybeConflict for surrealdb::Error {
55    fn is_transaction_conflict(&self) -> bool {
56        is_transaction_conflict(self)
57    }
58}
59
60/// Ends a transaction started by [`transaction!`](crate::transaction):
61/// commits it if `result` is `Ok`, cancels it otherwise.
62#[doc(hidden)]
63pub async fn finish<T, E: MaybeConflict>(
64    tx: Transaction,
65    result: std::result::Result<T, E>,
66) -> std::result::Result<T, E> {
67    match result {
68        Ok(value) => {
69            tx.commit().await?;
70            Ok(value)
71        }
72        Err(err) => {
73            let _ = tx.cancel().await;
74            Err(err)
75        }
76    }
77}
78
79/// Waits before retry `attempt` (2, 3, ...) of a conflicting transaction:
80/// exponential backoff with jitter, so the contenders spread out.
81#[doc(hidden)]
82pub async fn backoff(attempt: usize) {
83    let nanos = std::time::SystemTime::now()
84        .duration_since(std::time::UNIX_EPOCH)
85        .map_or(0, |d| d.subsec_nanos() as u64);
86    let base = 5u64 << attempt.min(6);
87    let millis = base + nanos % base;
88    tokio::time::sleep(std::time::Duration::from_millis(millis)).await;
89}
90
91/// Runs a block in a transaction: committed when it evaluates to `Ok`,
92/// cancelled when it evaluates to `Err` (`?` leaves the block, not the
93/// function). When the transaction conflicts with a concurrent one, the
94/// block runs again in a new transaction, up to
95/// [`TRANSACTION_ATTEMPTS`](database::TRANSACTION_ATTEMPTS) times.
96///
97/// ```ignore
98/// let user = slave_framework::transaction!(&db, |tx| {
99///     let user: Option<User> = tx.query("CREATE ...").await?.take(0)?;
100///     tx.query("CREATE ...").await?.check()?;
101///     Ok(user)
102/// })?;
103/// ```
104///
105/// The error type must implement [`MaybeConflict`](database::MaybeConflict).
106/// To make two transactions that read the same rows and write different ones
107/// conflict (e.g. "keep at least one owner"), have both write one shared
108/// record, such as the parent row.
109///
110/// A macro rather than a function taking an async closure: the futures of
111/// async closures are not provably `Send`, which MCP tools require.
112#[macro_export]
113macro_rules! transaction {
114    ($db:expr, |$tx:ident| $body:block) => {{
115        let db: &$crate::database::Db = $db;
116        let mut attempt = 1;
117        loop {
118            let result = match db.clone().begin().await {
119                Ok($tx) => {
120                    let result = async { $body }.await;
121                    $crate::database::finish($tx, result).await
122                }
123                Err(err) => Err(err.into()),
124            };
125            match result {
126                Err(err)
127                    if $crate::database::MaybeConflict::is_transaction_conflict(&err)
128                        && attempt < $crate::database::TRANSACTION_ATTEMPTS =>
129                {
130                    attempt += 1;
131                    $crate::database::backoff(attempt).await;
132                }
133                result => break result,
134            }
135        }
136    }};
137}
138
139/// How many times [`transaction!`](crate::transaction) runs a transaction
140/// that keeps failing with a conflict (see [`is_transaction_conflict`])
141/// before giving up.
142pub const TRANSACTION_ATTEMPTS: usize = 8;
143
144pub const URL_ENV: &str = "SURREAL_URL";
145pub const NAMESPACE_ENV: &str = "SURREAL_NS";
146pub const DATABASE_ENV: &str = "SURREAL_DB";
147pub const USER_ENV: &str = "SURREAL_USER";
148pub const PASSWORD_ENV: &str = "SURREAL_PASS";
149pub const AUTH_LEVEL_ENV: &str = "SURREAL_AUTH_LEVEL";
150
151/// The server used when `SURREAL_URL` is not set: `surreal start` on this machine.
152pub const DEFAULT_URL: &str = "ws://127.0.0.1:8000";
153
154/// Where a SurrealDB user is defined, which decides how to sign in.
155#[derive(Debug, Clone, Copy, PartialEq, Eq)]
156pub enum AuthLevel {
157    /// A root user (`surreal start --user ... --pass ...`).
158    Root,
159    /// A user defined on the namespace (`DEFINE USER ... ON NAMESPACE`).
160    Namespace,
161    /// A user defined on the database (`DEFINE USER ... ON DATABASE`), the
162    /// least privileged choice for production.
163    Database,
164}
165
166impl std::str::FromStr for AuthLevel {
167    type Err = String;
168
169    fn from_str(value: &str) -> std::result::Result<Self, String> {
170        match value.trim().to_ascii_lowercase().as_str() {
171            "root" => Ok(Self::Root),
172            "namespace" | "ns" => Ok(Self::Namespace),
173            "database" | "db" => Ok(Self::Database),
174            _ => Err("expected `root`, `namespace` or `database`".to_owned()),
175        }
176    }
177}
178
179/// Where to connect and how to sign in.
180#[derive(Debug, Clone)]
181pub struct DatabaseConfig {
182    /// `ws://host:port` or `wss://host` (SurrealDB Cloud).
183    pub url: String,
184    pub namespace: String,
185    pub database: String,
186    /// Signs in when set; a server started with `--unauthenticated` needs none.
187    pub username: Option<String>,
188    pub password: Option<String>,
189    pub auth_level: AuthLevel,
190}
191
192impl DatabaseConfig {
193    /// Reads `SURREAL_URL` (default [`DEFAULT_URL`]), `SURREAL_NS` and
194    /// `SURREAL_DB` (default `main`), `SURREAL_USER` and `SURREAL_PASS`, and
195    /// `SURREAL_AUTH_LEVEL` (`root`, `namespace` or `database`; default
196    /// `root`).
197    pub fn from_env() -> Result<Self> {
198        let non_empty = |key: &str| std::env::var(key).ok().filter(|v| !v.trim().is_empty());
199        Ok(Self {
200            url: crate::env::var_or(URL_ENV, DEFAULT_URL),
201            namespace: crate::env::var_or(NAMESPACE_ENV, "main"),
202            database: crate::env::var_or(DATABASE_ENV, "main"),
203            username: non_empty(USER_ENV),
204            password: non_empty(PASSWORD_ENV),
205            auth_level: crate::env::parse_or(AUTH_LEVEL_ENV, AuthLevel::Root)?,
206        })
207    }
208
209    /// The settings as environment variables, for child processes.
210    pub fn envs(&self) -> Vec<(&'static str, String)> {
211        let level = match self.auth_level {
212            AuthLevel::Root => "root",
213            AuthLevel::Namespace => "namespace",
214            AuthLevel::Database => "database",
215        };
216        let mut envs = vec![
217            (URL_ENV, self.url.clone()),
218            (NAMESPACE_ENV, self.namespace.clone()),
219            (DATABASE_ENV, self.database.clone()),
220            (AUTH_LEVEL_ENV, level.to_owned()),
221        ];
222        envs.extend(self.username.clone().map(|v| (USER_ENV, v)));
223        envs.extend(self.password.clone().map(|v| (PASSWORD_ENV, v)));
224        envs
225    }
226
227    /// `url ns/db`, for logs; never includes the password.
228    pub fn describe(&self) -> String {
229        format!("{} {}/{}", self.url, self.namespace, self.database)
230    }
231}
232
233/// Connects, signs in and selects the namespace and database.
234pub async fn connect(config: &DatabaseConfig) -> Result<Db> {
235    let db = surrealdb::engine::any::connect(config.url.as_str()).await?;
236    if let Some(username) = &config.username {
237        let password = config
238            .password
239            .clone()
240            .ok_or(Error::MissingDatabaseSetting(PASSWORD_ENV))?;
241        let username = username.clone();
242        match config.auth_level {
243            AuthLevel::Root => {
244                db.signin(Root { username, password }).await?;
245            }
246            AuthLevel::Namespace => {
247                db.signin(Namespace {
248                    namespace: config.namespace.clone(),
249                    username,
250                    password,
251                })
252                .await?;
253            }
254            AuthLevel::Database => {
255                db.signin(Database {
256                    namespace: config.namespace.clone(),
257                    database: config.database.clone(),
258                    username,
259                    password,
260                })
261                .await?;
262            }
263        }
264    }
265    db.use_ns(&config.namespace)
266        .use_db(&config.database)
267        .await?;
268    Ok(db)
269}
270
271/// Whether `err` is a transaction that lost a race with a concurrent one;
272/// running it again usually succeeds (up to [`TRANSACTION_ATTEMPTS`] times).
273pub fn is_transaction_conflict(err: &surrealdb::Error) -> bool {
274    matches!(err.query_details(), Some(QueryError::TransactionConflict))
275        || err.message().contains("can be retried")
276}
277
278/// Whether `err` is a write rejected by a `UNIQUE` index or an existing
279/// record ID.
280pub fn is_already_exists(err: &surrealdb::Error) -> bool {
281    err.is_already_exists() || err.message().contains("already contains")
282}
slave_framework/database.rs

slave_framework/
database.rs
