How to stream a query from sqlx using axum

[apps4uco]

Hi Id like to start a discussion about streaming from sqlx using hyper/axum, to show problems and solutions.

The first code is the working non streaming code to show how to generate ndjson.
This loads all the data into memory before streaming it out via http.

pub async fn products( params: Query<QueryParams>, Extension(pool): Extension<Pool<Postgres>>)
    -> impl IntoResponse {

    let products:Vec<Result<_, std::io::Error>>=query_as!(
        Product,
        "select * from products order by product_name")
            .map(|p| {
                let json=format!("{}\n",serde_json::to_string(&p).unwrap());
                Ok(Bytes::from(json)) })
            .fetch_all(&pool).await.unwrap();

    let stream = stream::iter(products);
    StreamBody::new(stream)
}

Working solution (without error handling)

pub async fn products( params: Query<QueryParams>, Extension(pool): Extension<Pool<Postgres>>)
                       -> impl IntoResponse
//Result<StreamBody<Iter<IntoIter<Result<Bytes,std::io::Error>>>>,AppError>
{
    let (mut sender,body)=hyper::body::Body::channel();

    tokio::spawn(async move {
        let mut products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>>=query_as!(
        Product,
        "select * from products order by product_name")
            .map(|p| {
                let json=format!("{}\n",serde_json::to_string(&p).unwrap());
                Bytes::from(json) })
            .fetch(&pool);

        info!("Start stream");
        while let Some(json)= products.try_next().await.unwrap() {
            sender.send_data(json).await.unwrap();
            // tokio::time::sleep(Duration::from_millis(100)).await;
        }
        info!("End stream");
    });

    body
}

Is this the most efficient solution?

Option 1: Loop can be changed to handle some errors, and handle end of the data, ie close json (not needed in this case as ndjson is line terminated)

 loop {
            let res= products.try_next().await;
            match res {
                Ok(Some(json)) => {
                    sender.send_data(json).await.unwrap();
                }
                Ok(None) => {
                    //End of data
                    info!("End of data");
                    break;
                }
                Err(err) => {
                    error!("Streaming Error {:?}",err);
                    //sender.send_trailers();
                    break;
                }
            }
            // tokio::time::sleep(Duration::from_millis(100)).await;
        }

Option 2: loop changed to handle some errors, slightly more ergonomic but does not handle end of data.

while let Some(res)= products.next().await {
            match res {
                Ok(json) => {
                    sender.send_data(json).await.unwrap();
                }
                Err(err) => {
                    error!("Streaming or Sql Error {:?}",err);
                    //sender.send_trailers();
                    break;
                }
            }
            // tokio::time::sleep(Duration::from_millis(100)).await;
        }

Previous attempts left to document what does not work.

If you have any suggestions as to how to make them work (if possible) or if you have another approach then please leave a comment.

Streaming Attempt 1

pub async fn products( params: Query<QueryParams>, Extension(pool): Extension<Pool<Postgres>>)
                       -> impl IntoResponse {
            let products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>> = query_as!(
            Product,
            "select * from products order by product_name")
                .map(|p| {
                    let json = format!("{}\n", serde_json::to_string(&p).unwrap());
                    Bytes::from(json)
                })
                .fetch(&pool);
            StreamBody::new(products)
}

This gives an error as the fetch requires a reference and that does not live long enough.

error[E0597]: `pool` does not live long enough
   --> src/northwind.rs:135:24
    |
127 |               let products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>> = query_as!(
    |  ___________________________________________________________________________________________________-
128 | |             Product,
129 | |             "select * from products order by product_name")
130 | |                 .map(|p| {
...   |
134 | |                 // .fetch(&mut connection);
135 | |                 .fetch(&pool);
    | |________________________^^^^^- cast requires that `pool` is borrowed for `'static`
    |                          |
    |                          borrowed value does not live long enough
...
139 |   }
    |   - `pool` dropped here while still borrowed

Streaming Attempt 2

pub async fn products( params: Query<QueryParams>, Extension(pool): Extension<Pool<Postgres>>)
                       -> impl IntoResponse {
        async move {
            let pool_clone=pool.clone();
            let products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>> = query_as!(
            Product,
            "select * from products order by product_name")
                .map(|p| {
                    let json = format!("{}\n", serde_json::to_string(&p).unwrap());
                    Bytes::from(json)
                })             
                .fetch(&pool_clone);
            // Ok(StreamBody::new(products))
            StreamBody::new(products)
        }.await
}

error[E0597]: `pool_clone` does not live long enough
   --> src/northwind.rs:136:24
    |
128 |               let products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>> = query_as!(
    |  ___________________________________________________________________________________________________-
129 | |             Product,
130 | |             "select * from products order by product_name")
131 | |                 .map(|p| {
...   |
135 | |                 // .fetch(&mut connection);
136 | |                 .fetch(&pool_clone);
    | |________________________^^^^^^^^^^^- cast requires that `pool_clone` is borrowed for `'static`
    |                          |
    |                          borrowed value does not live long enough
...
139 |           }.await
    |           - `pool_clone` dropped here while still borrowed

Streaming Attempt 3

pub async fn products( params: Query<QueryParams>, Extension(pool): Extension<Pool<Postgres>>)
                       -> impl IntoResponse {

        let mut connection=pool.acquire().await.unwrap();

        connection.transaction(|conn| async move {
            let products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>> = query_as!(
            Product,
            "select * from products order by product_name")
                .map(|p| {
                    let json = format!("{}\n", serde_json::to_string(&p).unwrap());
                    Bytes::from(json)
                })
                .fetch(conn);
            Ok(StreamBody::new(products))
        })
}

This looks like a compiler bug.

error[E0275]: overflow evaluating the requirement `&HashSet<_, _>: std::ops::Sub`
   --> src/northwind.rs:129:39
    |
129 |           connection.transaction(|conn| async move {
    |  _______________________________________^
130 | |             let products: Pin<Box<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>> + Send>> = query_as!(
131 | |             Product,
132 | |             "select * from products order by product_name")
...   |
139 | |             Ok(StreamBody::new(products))
140 | |         })
    | |_________^
    |
    = help: consider increasing the recursion limit by adding a `#![recursion_limit = "256"]` attribute to your crate (`server`)
    = note: required because of the requirements on the impl of `std::ops::Sub` for `&ordered_float::OrderedFloat<HashSet<_, _>>`
    = note: 126 redundant requirements hidden
    = note: required because of the requirements on the impl of `std::ops::Sub<&ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::Ord
...
dered_float::OrderedFloat<ordered_float::OrderedFloat<HashSet<_, _>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>`
    = note: required because of the requirements on the impl of `rulinalg::norm::MatrixMetric<_, rulinalg::matrix::MatrixSlice<'_, &ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat<ordered_float::OrderedFloat.....>` for `impl futures::Future`

Streaming Attempt 4

struct ProductStream {
    pool: Pool<Postgres>,
}

impl IntoResponse for ProductStream {
    type Body = StreamBody<Pin<Box<dyn futures::Stream<Item=Result<Bytes,axum::Error>> + Send>>>; //<dyn futures::Stream<Item=Result<Bytes, sqlx::Error>>;
    type BodyError = axum::Error;

    fn into_response(self) -> Response<Self::Body> {
        use futures::TryStreamExt;
        let stream //:Pin<Box<dyn futures::Stream<Item=Result<Bytes, (StatusCode,&str)>+Send>>>
            = query_as!(
            Product,
            "select * from products order by product_name")
            .map(|p| {
                let json = format!("{}\n", serde_json::to_string(&p).unwrap());
                Bytes::from(json)
            })
            .fetch(&self.pool)
               //.map_err(|e:sqlx::Error| (StatusCode::INTERNAL_SERVER_ERROR,String::from("Sql Error")))
            ;
        Response::new(StreamBody::new(products))
    }
}

pub async fn products( params: Query<QueryParams>, Extension(pool): Extension<Pool<Postgres>>)
                        -> impl IntoResponse {
    ProductStream{pool}
}

This looks like it might work but there is a problem with the types and errors.

error[E0308]: mismatched types
   --> src/northwind.rs:115:39
    |
115 |         Response::new(StreamBody::new(products))
    |                                       ^^^^^^^^ expected struct `Pin`, found fn item
    |
    = note: expected struct `Pin<Box<dyn futures::Stream<Item = Result<bytes::Bytes, axum::Error>> + std::marker::Send>>`
              found fn item `fn(axum::extract::Query<QueryParams>, axum::extract::Extension<sqlx::Pool<Postgres>>) -> impl futures::Future {products}`

It would be great if someone could help out and edit the code to get a working example.

[davidpdrsn]

If sqlx's streams borrow the connection (or connection pool) then its not compatible with axum's StreamBody since it requires a 'static stream, meaning it must be fully owned and not contain borrows.

You're also trying to stream sql rows which you cannot stream as raw HTTP bodies (unless sqlx's streams are already of Vec<u8>s). You can only stream something that can be converted into Bytes. This is also shown in the trait bounds of StreamBody::new.

I'm not familiar enough with sqlx to know how to make it work (or if its even possible). I think you'll get better help asking them.

[apps4uco]

Hi thanks for your help,

I found a way to do it using hyper::body::Body::channel() and spawning a tokio task. See second code example. TODO: add error handling, does anyone have any suggestions?

[davidpdrsn]

You mean handling errors after you've sent the response and are streaming the body? Thats pretty tricky. I would recommend looking at what gRPC does which uses a trailing http2 header to indicate success or failure.

[belovaf]

I find the solution.
You can use Box::leak to make state static:

struct AppState {
    db_pool: PgPool,
}

impl <'a> FromRef<&'a AppState> for &'a PgPool {
    fn from_ref(input: &&'a AppState) -> Self {
        &input.db_pool
    }
}

#[tokio::main]
async fn main() {
    // init db, http client, etc...

    let state: &'static AppState = Box::leak(Box::new(AppState {
        db_pool,
        http_client,
    }));
    
    let app = Router::new()
        .route("/", get(handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn handler(State(db_pool): State<&'static PgPool>) -> impl IntoResponse {
    let stream = sqlx::query_as!(Row, r#"SELECT col from my_table"#).fetch(db_pool);
    StreamBody::new(stream)
}

struct Row {
    col: String
}

impl From<Row> for Bytes {
    fn from(value: Row) -> Self {
        Bytes::from(value.col)
    }
}

[cbeck88]

i would not do this, because if you want to use axum_test to test your routes, you may have weird issues. for any test that uses a route that takes &'static you would have to leak an AppState. But you may have a lot of tests. Eventually it could result in things like, test db's not cleaned up, other resources not released.