Pagination in Async Diesel

In the Rust ecosystem, there are quite a few ways to approach interacting with Postgres. One of the older and more established libraries is diesel, which effectively gives you a Rust-wrapped path to constructing and executing queries. If it helps, you can think of it as a "light" ORM.

One of the bigger questions over the years has been whether diesel not being "async by default" is an issue. You could always wrap your queries in tasks and just await those, and that's probably fine - in fact, if I understand correctly, this is what crates.io does. However, diesel-async also exists for those who want to have a "pure" approach to the async interaction pattern. It's a smaller crate that effectively just requires importing a few extra traits to "override" the core diesel traits, and changing your connection type to be an asynchronous one. Overall, I've found it works pretty well.

One thing I did run into recently was needing some pagination for an admin interface I was throwing together. diesel actually has an example for this that works well if you're using synchronous calls or wrapping in tasks, but if you're using diesel-async, you'll need to tweak some things to make it work. I figured I'd share my changes here in case anyone needs them; yes, these could be on GitHub, but I'm in a weird phase of re-evaluating my usage of the platform, so I'm dumping it here.

You're smart, you can figure it out. The only real changes needed were importing some diesel-async items and tweaking the bounds on load_and_count_pages to accomodate the async differences. With this, you should be able to just call .paginate(page) on your queries and have mostly automatic paging. Enjoy.

use diesel::pg::Pg;
use diesel::prelude::*;
use diesel::query_builder::*;
use diesel::sql_types::BigInt;

// Import these, since we'll need them.
use diesel_async::AsyncPgConnection;
use diesel_async::methods::LoadQuery;

pub trait Paginate: Sized {
    fn paginate(self, page: i64) -> Paginated<Self>;
}

impl<T> Paginate for T {
    fn paginate(self, page: i64) -> Paginated<Self> {
        Paginated {
            query: self,
            per_page: DEFAULT_PER_PAGE,
            page,
            offset: (page - 1) * DEFAULT_PER_PAGE,
        }
    }
}

const DEFAULT_PER_PAGE: i64 = 10;

#[derive(Debug, Clone, Copy, QueryId)]
pub struct Paginated<T> {
    query: T,
    page: i64,
    per_page: i64,
    offset: i64,
}

impl<T> Paginated<T> {
    pub fn per_page(self, per_page: i64) -> Self {
        Paginated {
            per_page,
            offset: (self.page - 1) * per_page,
            ..self
        }
    }

    // Mark the bounds on `T` and `U` to ensure the async calls compile correctly,
    // and have `results` go through the async `RunQueryDsl`.
    pub async fn load_and_count_pages<'a, U>(
        self,
        conn: &mut AsyncPgConnection
    ) -> QueryResult<(Vec<U>, i64)>
    where
        T: 'a,
        U: Send + 'a,
        Self: LoadQuery<'a, AsyncPgConnection, (U, i64)>,
    {
        let per_page = self.per_page;
        let results = diesel_async::RunQueryDsl::load::<(U, i64)>(self, conn).await?;
        let total = results.first().map(|x| x.1).unwrap_or(0);
        let records = results.into_iter().map(|x| x.0).collect();
        let total_pages = (total as f64 / per_page as f64).ceil() as i64;
        Ok((records, total_pages))
    }
}

impl<T: Query> Query for Paginated<T> {
    type SqlType = (T::SqlType, BigInt);
}

impl<T> RunQueryDsl<PgConnection> for Paginated<T> {}

impl<T> QueryFragment<Pg> for Paginated<T>
where
    T: QueryFragment<Pg>,
{
    fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
        out.push_sql("SELECT *, COUNT(*) OVER () FROM (");
        self.query.walk_ast(out.reborrow())?;
        out.push_sql(") t LIMIT ");
        out.push_bind_param::<BigInt, _>(&self.per_page)?;
        out.push_sql(" OFFSET ");
        out.push_bind_param::<BigInt, _>(&self.offset)?;
        Ok(())
    }
}