In rust, any way to return `rows_affected` from an `sqlx::query_with` along with the rows themselves?

Question

The sqlx documentation says the following:

The execute query finalizer returns the number of affected rows, if any, and drops all received results. In addition, there are fetch, fetch_one, fetch_optional, and fetch_all to receive results.

The Query type returned from sqlx::query will return Row<'conn> from the database. Column values can be accessed by ordinal or by name with row.get(). As the Row retains an immutable borrow on the connection, only one Row may exist at a time.

Is there any way get both the rows_affected and the rows themselves?

I tried running the same query twice in a transaction that I rollback intentionally, but I am already in the middle of a transaction when I receive the sql statement, and I cannot run two transactions in parallel.

Is there a way to do it without running it twice?

I am using rust, the latest sqlx and postgres as the database.

Answer 1

You'll also need the futures crate:

[dependencies]
futures = { version = "0.3.28", default-features = false }
sqlx = { version = "0.7.2", features = ["runtime-tokio", "postgres"] }

Then you can do shenanigans like this:

use futures::stream::StreamExt;
use sqlx::{postgres::PgPoolOptions, Either, PgPool, Row};

#[tokio::main]
async fn main() {
    # ...

    do_query(
        &pool,
        "INSERT INTO todos (description) VALUES ('One'), ('Two') RETURNING id",
    )
    .await;
    do_query(&pool, "DELETE FROM todos").await;
}

async fn do_query(pool: &PgPool, query: &str) {
    let prefix = query.split_whitespace().next().unwrap();

    let mut results = sqlx::query(query).fetch_many(pool);

    while let Some(result) = results.next().await {
        let either = result.unwrap();
        match either {
            Either::Left(res) => {
                let num = res.rows_affected();
                println!("[{prefix}] affected {num} rows");
            }
            Either::Right(row) => {
                let num = row.len();
                println!("[{prefix}] fetched {num} rows");
            }
        }
    }
}

This will return

[INSERT] fetched 1 rows
[INSERT] fetched 1 rows
[INSERT] affected 2 rows
[DELETE] affected 2 rows