How to write a method that can take either a Transaction or a PgPool?

[ecton]

I'm trying to write code like this:

#[derive(Debug, Clone, FromRow, Serialize, Deserialize)]
pub struct Account {
    pub id: i64,
...
}

impl Account {
    pub async fn save<'e, E: sqlx::Executor<'e, Database = sqlx::Postgres>>(
        &mut self,
        executor: E,
    ) -> sqlx::Result<()> {
        let row = sqlx::query!(
            "INSERT INTO accounts (username, password_hash, display_name) VALUES ($1, $2, $3) RETURNING id, created_at",
            &self.username,
            &self.password_hash,
            self.display_name.as_ref(),
        )
        .fetch_one(executor)
        .await?;
        self.id = row.id;
        self.created_at = row.created_at;
    }
}

This works well on its own, but the issue is when I want to use it in a method like this:

pub async fn setup_test_account(
    transaction: &mut Transaction<'_, Postgres>,
) -> anyhow::Result<(Account, Session)> {
    let mut account = Account::new(TEST_ACCOUNT_USERNAME, TEST_ACCOUNT_PASSWORD)?;
    account.save(transaction).await?;

    let session = Session::new(&account, None, transaction).await?;

    Ok((account, session))
}

Regardless of if I try to define setup_test_account as taking a generic sqlx::Executor, I cannot figure out how to make save() operate both inside and outside of a transaction. If I hardcode save() to always use a Transaction, that's not a big deal in-and-of-itself, but I'm using this pattern all across my data access layer, including methods such as load(). I don't want to create transactions for calling load(), and I don't want to have two method definitions -- one for when calling within a transaction and one for when not.

The issue stems from the exclusive reference to the transaction being passed by-value into save(). This happens when you call sqlx::query!([...]).fetch_one(&mut pool) because the mutable reference is only acquired for that single statement, and further in the method, when you use it again, you'll create a new reference.

I've tried a lot, but ultimately I've not found an approach that works. Does anyone have any ideas?

Any help is very appreciated.

[pfernie]

It is unclear from your description if you are trying to keep the signature of setup_test_account (taking &mut Transaction), and allowing for setup_test_account to call onto save and then Session::new. If that is what you're looking to do, re-borrowing transaction will accomplish this:

pub async fn setup_test_account_simple(
    transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
) -> Result<(Account, Session)> {
    let mut account = Account::new(TEST_ACCOUNT_USERNAME, TEST_ACCOUNT_PASSWORD);
    account.save(&mut *transaction).await?;

    let session = Session::new(&account, None, &mut *transaction).await?;

    Ok((account, session))
}

However, if you are more generally trying to be able to generically pass "executors" down some arbitrary set of calls, then I believe you need to reach for Acquire.

pub async fn setup_test_account<'e, A: sqlx::Acquire<'e, Database = sqlx::Postgres>>(
    acquire_from: A,
) -> Result<(Account, Session)> {
    let mut connection = acquire_from.acquire().await?;

    let mut account = Account::new(TEST_ACCOUNT_USERNAME, TEST_ACCOUNT_PASSWORD);
    account.save(&mut *connection).await?;`

    let session = Session::new(&account, None, &mut *connection).await?;

    Ok((account, session))
}

Where the &mut * on each use of connection are taking advantage on Acquire's bounds on type Connection: Deref<...> + DerefMut.

Depending on how deeply one might need to keep this generic, the "actual" executing fns can take Acquire as well:

impl Account {
..[SNIP]..
    pub async fn frobnicate<'e, A: sqlx::Acquire<'e, Database = sqlx::Postgres>>(
        &mut self,
        acquire_from: A,
    ) -> sqlx::Result<()> {
        let mut connection = acquire_from.acquire().await?;
        let row = sqlx::query!("SELECT COUNT(*) AS count FROM accounts",)
            .fetch_one(&mut *connection)
            .await?;
        dbg!(row.count);
        Ok(())
    }
}

Would allow:

pub async fn setup_test_account<'e, A: sqlx::Acquire<'e, Database = sqlx::Postgres>>(
    acquire_from: A,
) -> Result<(Account, Session)> {
    let mut connection = acquire_from.acquire().await?;

    let mut account = Account::new(TEST_ACCOUNT_USERNAME, TEST_ACCOUNT_PASSWORD);
    account.save(&mut *connection).await?;

    account.frobnicate(&mut *connection).await?; // <-- This works as well

    let session = Session::new(&account, None, &mut *connection).await?;

    Ok((account, session))
}

With these definitions, the following works (my test code is slightly changed from your example, but hopefully illustrative):

    {
        let mut account = Account::new("Gene".into());
        account.save(&pool).await?;
        dbg!(&account);

        dbg!(setup_test_account(&pool, "Wolfe".into()).await?);
    }

    {
        let mut transaction = pool.begin().await?;
        dbg!(setup_test_account(&mut *transaction, "Zaphod".into()).await?);
        dbg!(setup_test_account(&mut *transaction, "Douglas".into()).await?);
        transaction.commit().await?;
    }

    {
        let mut transaction = pool.begin().await?;
        dbg!(setup_test_account(&mut *transaction, "Beeblebrox".into()).await?);
        dbg!(setup_test_account(&mut *transaction, "Adams".into()).await?);
        // we expect this transaction to be aborted
        //transaction.commit().await?;
    }

[jplatte]

A value of type E where E: Exectuor<'e, DB> can only be used once. A &mut E with the same constraint would work if Rust had lazy normalization (see #588), but it doesn't. There are two options that you can choose from:

• Use &mut PgConnection<'_> instead of &mut Transaction<'_, Postgres>
  • &mut Transaction<'_, Postgres> derefs to &mut PgConnection, making it easy to use a transaction
  • When you have a pool, you need to use .acquire().await (or .begin().await) to get a connection or transaction to pass to setup_test_account
• Use &mut A where A: Acquire<'c, Database = Postgres>
  • Then you need to .acquire().await inside setup_test_account to get a connection

[ecton]

Thank you both @jplatte and @pfernie. I somehow completely missed the link that Acquire could be used this way. Apologies for the unclear question, I realize in hindsight my question could be better summed up as "how to write a method that can accept either a transaction or a pool that can make two sqlx macro-based queries".

Both of your answers give ways to approach it that I just hadn't connected the dots on yet.