Quick way to emulate a database in Rust

Recently I was working with shuttle.rs, a Backend-as-a-Service platform built on Rust, building a simple web app with Axum and Postgres. There was a GET handler that reads from and a POST handler that writes to the database, which is provisioned by Shuttle:

async fn axum(
    #[shuttle_shared_db::Postgres] pool: PgPool,
) -> shuttle_axum::ShuttleAxum {
    let router =
        Router::new().nest_service(
            "/posts/:id",
            get(read_post)
        ).nest_service(
            "/new-post",
            post(create_post)
        );

    Ok(router.into())
}

async fn read_post(Path(id): Path<String>) -> Html<Bytes> {
    ...
}

async fn create_post(body: Bytes) -> impl IntoResponse {
    ...
}

When I was doing the integration between the routing on the server and the frontend, I wanted to test these handlers, but I dreaded the prospect of spinning up a whole Postgres instance and doing all the schema migrations that entails. Moreover, I would need to modify the handlers read_post and create_post, passing in the PgPool from sqlx by using closures, of which is I am not very fond.

All I need is a global mutable key-value store. For testing purposes, it does not need to persist, i.e., it can be an in-memory struct. “Easy enough, just use a HashMap in axum, which even avoids global state,” I said:

async fn axum(...) {
    let mut store = HashMap::new();
        let router =
        Router::new().nest_service(
            "/posts/:id",
            get(|path| read_post(&store, path))
        ).nest_service(
            "/new-post",
            post(|body| create_post(&mut store, body))
        );

    Ok(router.into())
}

Assuming I have modified the handlers accordingly, the code fails with some unsatisfied trait bound on the closures, which is a little Axum-specific. Other than that, the borrow checker is unhappy as well: store is owned by the axum function and borrowed by the handlers, but the handlers may live longer than axum (indeed, the handlers can run any time in the future, long after axum returns).

What else? We have to resort to static variables, which is guaranteed to exist throughout the program’s life. However, mutating a static variable requires unsafe code,¹ which I would like to avoid if possible. To that end, we can wrap the HashMap around a Mutex. Usually, people turn such variables into singletons using some libraries like lazy-static or once-cell, but as of Rust v1.72.0, we can just use std::sync::OnceLock. In short,

fn db() -> &'static Mutex<HashMap<String, Bytes>> {
    static HASHMAP: OnceLock<Mutex<HashMap<String, Bytes>>> = OnceLock::new();
    HASHMAP.get_or_init(|| Mutex::new(HashMap::new()))
}

async fn read_post(Path(id): Path<String>) -> Html<Bytes> {
    let db = db().lock().expect("Failed to get database");
    let content = db.get(&id);
    ...
}

async fn create_post(body: Bytes) -> impl IntoResponse {
    let mut db = db().lock().expect("Failed to get database");
    ...
}

And there you go. A simple and quick way to get a global mutable singleton key-value in-memory store.²