Queries and Handles

Define entities and queries, run them through capability-typed Db handles, route reads to a replica, and inject the handles into your components.

This page builds the data access of a small user service on PostgreSQL. It defines queries, builds a primary handle and a replica handle, sends each query to the right one, and registers the handles in the ServiceBuilder. The concepts behind it (entities, capabilities, effects) are explained in the overview.

Setup

Cargo.toml

[dependencies]
wakuwaku = { version = "1.0.0-rc3", features = ["sqlx-pg"] }
tokio = { version = "1", features = ["full"] }

sqlx-pg enables wakuwaku::entity with the PostgreSQL data source, and re-exports sqlx as wakuwaku::deps::sqlx.

Define the entities

src/entities.rs

use wakuwaku::entity::Entity;

/// The primary key of [`User`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct UserId(pub i64);

/// The `users` table.
pub struct User;
impl Entity for User {}

/// The primary key of [`Post`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct PostId(pub i64);

/// The `posts` table.
pub struct Post;
impl Entity for Post {}

Give each primary key its own type

A bare i64 doesn't say which table it identifies. PostsByAuthor { author } accepts a post id as readily as a user id, and swapping id and version in EditPostTitle compiles. A newtype per primary key turns these mix-ups into type errors, and it documents foreign keys: author: UserId says the field points at users, and "Go to Type Definition" on author anywhere it is used takes you to UserId, right next to the User entity. The wrapper costs nothing at runtime. Unwrap it with .0 where you bind it, and wrap it where you read it.

If you also depend on sqlx directly, at the version wakuwaku uses and with its derive feature, #[derive(sqlx::Type)] with #[sqlx(transparent)] lets you bind and decode the newtype itself.

Define a query

A query implements two traits:

Query says what it touches: its Effect. This part doesn't depend on where the query runs, so capability checks and hooks work the same for every data source.
Execute<S> says how it runs on the data source S: its Output, its Error, and the code that runs on a connection. A query written in portable SQL may implement it for several sources.

src/queries.rs

use crate::entities::{Post, PostId, User, UserId};
use wakuwaku::deps::sqlx::{self, PgConnection};
use wakuwaku::entity::sqlx::backend::PgSource;
use wakuwaku::entity::{Execute, Query, Read, Write};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UserRow {
    pub id: UserId,
    pub name: String,
}

/// Load one user by id.
pub struct FindUser {
    pub id: UserId,
}

impl Query for FindUser {
    type Effect = Read<User>;
}

impl Execute<PgSource> for FindUser {
    type Output = Option<UserRow>;
    type Error = sqlx::Error;

    async fn execute(self, conn: &mut PgConnection) -> Result<Option<UserRow>, sqlx::Error> {
        let row: Option<(i64, String)> = sqlx::query_as("SELECT id, name FROM users WHERE id = $1")
            .bind(self.id.0)
            .fetch_optional(conn)
            .await?;
        Ok(row.map(|(id, name)| UserRow { id: UserId(id), name }))
    }
}

execute receives a plain &mut PgConnection and returns the query's Error. A query that can only fail the way the database fails uses the data source's error, sqlx::Error for sqlx sources; a query that can also refuse on its own declares its own error type. When the query runs on its own, the connection comes from the pool in autocommit mode. Inside a transaction, it is the transaction's connection. The same query works in both places.

A query that touches several entities lists them all. A tuple is the union of its effects:

src/queries.rs

/// Titles of the posts of an active author.
pub struct PostsByAuthor {
    pub author: UserId,
}

impl Query for PostsByAuthor {
    // Reads two tables: a tuple is the union of its effects.
    type Effect = (Read<User>, Read<Post>);
}

impl Execute<PgSource> for PostsByAuthor {
    type Output = Vec<String>;
    type Error = sqlx::Error;

    async fn execute(self, conn: &mut PgConnection) -> Result<Vec<String>, sqlx::Error> {
        sqlx::query_scalar(
            "SELECT p.title FROM posts p JOIN users u ON u.id = p.author_id \
             WHERE u.id = $1 AND u.active ORDER BY p.created_at DESC",
        )
        .bind(self.author.0)
        .fetch_all(conn)
        .await
    }
}

A write declares Write<E>. A capability that can write E can also read it, so Write<User> is enough for an UPDATE that also filters on users:

src/queries.rs

/// Change a user's display name.
pub struct RenameUser {
    pub id: UserId,
    pub name: String,
}

impl Query for RenameUser {
    type Effect = Write<User>;
}

impl Execute<PgSource> for RenameUser {
    /// Rows affected.
    type Output = u64;
    type Error = sqlx::Error;

    async fn execute(self, conn: &mut PgConnection) -> Result<u64, sqlx::Error> {
        let done = sqlx::query("UPDATE users SET name = $2 WHERE id = $1")
            .bind(self.id.0)
            .bind(self.name)
            .execute(conn)
            .await?;
        Ok(done.rows_affected())
    }
}

Declare honestly

The compiler checks the effect you declare, not the SQL you write. A query declared as Read<User> that runs an UPDATE still compiles. Treat the Effect line like a function signature and review it together with the statement.

Data sources

A DataSource is anything that can run queries: a connection pool, a multiplexed client. It fixes the connection type that queries receive and its own error type, which every query's Error must be convertible from:

pub trait DataSource: Send + Sync + Sized {
    type Connection: Send;
    type Error: Send;

    fn run<Q: Execute<Self>>(
        &self,
        query: Q,
    ) -> impl Future<Output = Result<Q::Output, Q::Error>> + Send;
}

run returns the query's own result. If the source fails before the query starts, for example because no connection could be acquired, that error is converted into Q::Error through From<Self::Error>.

You rarely implement it yourself. wakuwaku provides one per backend:

Source	Feature	Connection	Source error
`entity::sqlx::backend::PgSource`	`sqlx-pg`	`PgConnection`	`sqlx::Error`
`entity::sqlx::backend::MySqlSource`	`sqlx-mysql`	`MySqlConnection`	`sqlx::Error`
`entity::redis::RedisSource`	`redis`	see Redis	`redis::RedisError`
`entity::meilisearch::MeiliSource`	`meilisearch`	see Meilisearch	`meilisearch_sdk::errors::Error`

PgSource is an alias for SqlxSource<Postgres>. Each query run on its own acquires a connection from the pool.

Build the handles

A Db<S, C> owns a data source S and has a capability C. The capability is only a type parameter, so you choose it when you build the handle. A service with a primary and a read replica builds two handles:

src/wiring.rs

use crate::entities::{PostId, UserId};
use crate::queries::{
    DeactivateUsers, DormantUsers, EditPostError, EditPostTitle, FindUser, PostsByAuthor, RenameUser, UserRow,
};
use std::time::Duration;
use wakuwaku::deps::kanau::processor::Processor;
use wakuwaku::deps::sqlx::{self, postgres::PgPoolOptions};
use wakuwaku::entity::sqlx::backend::{PgSource, SqlxSource};
use wakuwaku::entity::{Db, ReadOnly, ReadWrite};
use wakuwaku::{CronJob, CronJobInterval, CronTick, ServiceBuilder};

/// The two handles of the service: writes go to the primary,
/// reads that tolerate replication lag go to the replica.
#[derive(Clone)]
pub struct Databases {
    pub primary: Db<PgSource, ReadWrite>,
    pub replica: Db<PgSource, ReadOnly>,
}

impl Databases {
    pub async fn connect(primary_url: &str, replica_url: &str) -> Result<Self, sqlx::Error> {
        let primary = PgPoolOptions::new().max_connections(10).connect(primary_url).await?;
        let replica = PgPoolOptions::new().max_connections(20).connect(replica_url).await?;
        Ok(Self {
            primary: Db::new(SqlxSource::new(primary)),
            replica: Db::new(SqlxSource::new(replica)),
        })
    }
}

The type is the promise

Db::new can't tell which server a pool is connected to. Building a Db<PgSource, ReadWrite> from the replica pool compiles, and it would then send writes to the replica. Only build a writing capability from a source connected to a writable database, and keep this wiring in one place.

Once a pool is wrapped, you can't get it back out of the handle. That is what makes the capability sound. Run migrations and other setup on the pool before calling SqlxSource::new. Cloning a Db clones its source, and for sqlx that means cloning the pool handle, which is cheap.

Run a query

Db implements kanau's Processor. Bring the trait into scope and call process:

src/wiring.rs

pub async fn run_one(dbs: &Databases) -> Result<(), sqlx::Error> {
    let user: Option<UserRow> = dbs.replica.process(FindUser { id: UserId(7) }).await?;
    let renamed: u64 = dbs.primary.process(RenameUser { id: UserId(7), name: "Haruki".into() }).await?;
    println!("{user:?}, {renamed} row(s) updated");
    Ok(())
}

process returns the query's Output, or the query's Error.

Keep handlers generic over the handle

Handlers shouldn't name Db<PgSource, ReadOnly>. They should list the queries they need as Processor bounds:

src/wiring.rs

pub struct Profile {
    pub user: UserRow,
    pub posts: Vec<String>,
}

/// Works with any handle that can run both queries.
pub async fn load_profile<D>(db: &D, id: UserId) -> Result<Option<Profile>, sqlx::Error>
where
    D: Processor<FindUser, Output = Option<UserRow>, Error = sqlx::Error>
        + Processor<PostsByAuthor, Output = Vec<String>, Error = sqlx::Error>,
{
    let Some(user) = db.process(FindUser { id }).await? else {
        return Ok(None);
    };
    let posts = db.process(PostsByAuthor { author: id }).await?;
    Ok(Some(Profile { user, posts }))
}

pub async fn rename_user<D>(db: &D, id: UserId, name: String) -> Result<bool, sqlx::Error>
where
    D: Processor<RenameUser, Output = u64, Error = sqlx::Error>,
{
    Ok(db.process(RenameUser { id, name }).await? == 1)
}

The capability is then checked once, where the handle is passed in. This also means:

A handler accepts any processor of its queries: a Db, a Db wrapped in write hooks, or a fake in a unit test.
With several Processor bounds, pin Output and Error in each bound as above, so that db.process(...) has an unambiguous result type.

A query's own error type

Some queries can fail for reasons of their own, not only because the database failed. An update guarded by an optimistic lock finds that someone else changed the row first; a withdrawal would overdraw the account. Such a refusal is a failure of the operation, so it belongs in the query's Error, not in its Output.

Putting it in the Output (Output = bool, or Output = Result<Post, Conflict>) works, but every caller then has to check the result a second time after the ?, and a caller that forgets compiles fine and treats the refusal as success. With its own error type, the refusal travels like any other error: ? propagates it and a match handles it.

The only requirement is type Error: From<S::Error>, so that the data source's errors fit into it. Here a posts row carries a version column, and an edit applies only to the version the client last read:

src/queries.rs

/// Why an edit of a post failed.
#[derive(Debug)]
pub enum EditPostError {
    /// The post changed, or was deleted, since `version` was read.
    Stale,
    /// The database failed.
    Database(sqlx::Error),
}

impl From<sqlx::Error> for EditPostError {
    fn from(error: sqlx::Error) -> Self {
        Self::Database(error)
    }
}

/// Change a post's title, if nobody changed the post since `version` was read.
pub struct EditPostTitle {
    pub id: PostId,
    pub version: i64,
    pub title: String,
}

impl Query for EditPostTitle {
    type Effect = Write<Post>;
}

impl Execute<PgSource> for EditPostTitle {
    /// The new version of the post.
    type Output = i64;
    type Error = EditPostError;

    async fn execute(self, conn: &mut PgConnection) -> Result<i64, EditPostError> {
        let version: Option<i64> = sqlx::query_scalar(
            "UPDATE posts SET title = $3, version = version + 1 \
             WHERE id = $1 AND version = $2 RETURNING version",
        )
        .bind(self.id.0)
        .bind(self.version)
        .bind(self.title)
        .fetch_optional(conn)
        .await?; // sqlx::Error -> EditPostError::Database
        version.ok_or(EditPostError::Stale)
    }
}

Db::process returns Result<i64, EditPostError>, so the caller decides what each failure means in one place:

src/wiring.rs

pub async fn save_title(dbs: &Databases, id: PostId, version: i64, title: String) -> Result<Option<i64>, sqlx::Error> {
    match dbs.primary.process(EditPostTitle { id, version, title }).await {
        Ok(new_version) => Ok(Some(new_version)),
        // Someone else won: tell the client to reload and retry.
        Err(EditPostError::Stale) => Ok(None),
        Err(EditPostError::Database(error)) => Err(error),
    }
}

A function that returns EditPostError can also ? queries whose error is sqlx::Error, through the same From impl. Hooks don't change any of this: a hooked handle returns the query's error unchanged.

Inside a transaction

A query run inside a transaction must have an error that implements AsSqlxError. It tells the transaction whether the error came from the database, which may have aborted the transaction, or is the query's own, which never does. sqlx::Error implements it; for your own error type, return the wrapped sqlx::Error:

src/queries.rs

use wakuwaku::entity::sqlx::transaction::AsSqlxError;

impl AsSqlxError for EditPostError {
    fn as_sqlx_error(&self) -> Option<&sqlx::Error> {
        match self {
            Self::Database(error) => Some(error),
            Self::Stale => None,
        }
    }
}

See errors inside a transaction.

Route reads to the replica and writes to the primary

At the call site, pick the handle:

src/wiring.rs

pub async fn handle_requests(dbs: &Databases) -> Result<(), sqlx::Error> {
    // A profile page: a slightly stale read is fine.
    let _profile = load_profile(&dbs.replica, UserId(7)).await?;

    // Writes go to the primary.
    rename_user(&dbs.primary, UserId(7), "Haruki".into()).await?;

    // Read your own write: the primary can read too.
    let _fresh = load_profile(&dbs.primary, UserId(7)).await?;
    Ok(())
}

ReadWrite covers every read, so the primary handle runs read queries too. Use it when a read must see a write that was just made, because a replica can lag behind the primary.

When a write reaches a read-only handle

Pass the replica to a handler that writes:

pub async fn oops(dbs: &Databases) -> Result<bool, sqlx::Error> {
    rename_user(&dbs.replica, UserId(7), "Haruki".into()).await
}

The program doesn't compile:

error[E0277]: the trait bound `ReadOnly: CanWrite<User>` is not satisfied
  --> src/fails.rs:18:17
   |
18 |     rename_user(&dbs.replica, UserId(7), "Haruki".into()).await
   |     ----------- ^^^^^^^^^^^^ the trait `CanWrite<User>` is not implemented for `ReadOnly`
   |     |
   |     required by a bound introduced by this call
   |
   = note: required for `ReadOnly` to implement `Covers<wakuwaku::wakuwaku_entity::Write<User>>`
   = note: required for `Db<SqlxSource<Postgres>, ReadOnly>` to implement `Processor<RenameUser>`
note: required by a bound in `rename_user`

Read the notes from the bottom up. rename_user needs Processor<RenameUser>. Db provides it only if its capability covers Write<User>, and ReadOnly has no CanWrite<User>. Calling dbs.replica.process(RenameUser { .. }) directly fails with the same error.

If a component should write only some entities, give it a custom capability instead of ReadWrite.

Inject the handles into components

Handles are ordinary values, so they go into the ServiceBuilder like any other shared dependency. Db<PgSource, ReadOnly> and Db<PgSource, ReadWrite> are different types, so register both and let each component take the ones it needs.

This cron job finds dormant users on the replica and deactivates them on the primary:

src/wiring.rs

/// Deactivate users that haven't signed in for a year.
pub struct DeactivateDormantUsers {
    replica: Db<PgSource, ReadOnly>,
    primary: Db<PgSource, ReadWrite>,
}

impl Processor<CronTick> for DeactivateDormantUsers {
    type Output = ();
    type Error = sqlx::Error;

    async fn process(&self, _tick: CronTick) -> Result<(), sqlx::Error> {
        let ids = self.replica.process(DormantUsers { days: 365 }).await?;
        if !ids.is_empty() {
            self.primary.process(DeactivateUsers { ids }).await?;
        }
        Ok(())
    }
}

impl CronJob for DeactivateDormantUsers {
    const INTERVAL: CronJobInterval = CronJobInterval::AfterFinish(Duration::from_secs(60 * 60));
}

DormantUsers (Read<User>, outputs Vec<UserId>) and DeactivateUsers (Write<User>, takes ids: Vec<UserId>) are defined like the queries above. Register the handles, then build the job from both of them:

src/wiring.rs

pub async fn start() -> Result<(), sqlx::Error> {
    let dbs = Databases::connect(
        "postgres://app@db-primary/app",
        "postgres://app@db-replica/app",
    )
    .await?;

    let (services, cron) = ServiceBuilder::new(dbs.primary.clone())
        .push(dbs.replica.clone())
        .compose_cron_job(|s| DeactivateDormantUsers {
            replica: s.provide::<Db<PgSource, ReadOnly>, _>().clone(),
            primary: s.provide::<Db<PgSource, ReadWrite>, _>().clone(),
        })
        .start();

    // Everything else gets the handles from the builder too.
    let replica: &Db<PgSource, ReadOnly> = services.provide::<Db<PgSource, ReadOnly>, _>();
    let _profile = load_profile(replica, UserId(7)).await?;
    handle_requests(&dbs).await?;

    cron.shutdown().await;
    Ok(())
}

compose_cron_job lends the closure the builder, so a component can take any number of registered handles; it needs no ServiceCreation impl. A component that needs a single handle can implement ServiceCreation with that handle as its Dep and be registered with cron_job. gRPC services and event consumers have the same pair of methods; see ServiceBuilder. Cloning a handle only clones the pool handle, so every component can hold its own copy.

Queries and Handles

Setup

Define the entities

Define a query

Data sources

Build the handles

Run a query

Keep handlers generic over the handle

A query's own error type

Route reads to the replica and writes to the primary

When a write reaches a read-only handle

Inject the handles into components

Next steps

Write hooks

sqlx

Transactions

On this page

Queries and Handles

SetupCopy Anchor Link

Define the entitiesCopy Anchor Link

Define a queryCopy Anchor Link

Data sourcesCopy Anchor Link

Build the handlesCopy Anchor Link

Run a queryCopy Anchor Link

Keep handlers generic over the handleCopy Anchor Link

A query's own error typeCopy Anchor Link

Route reads to the replica and writes to the primaryCopy Anchor Link

When a write reaches a read-only handleCopy Anchor Link

Inject the handles into componentsCopy Anchor Link

Why are queries values instead of repository methods?

Why is the capability a type parameter and not a runtime flag?

Next stepsCopy Anchor Link

Write hooks

sqlx

Transactions

On this page

Setup

Define the entities

Define a query

Data sources

Build the handles

Run a query

Keep handlers generic over the handle

A query's own error type

Route reads to the replica and writes to the primary

When a write reaches a read-only handle

Inject the handles into components

Next steps