systemprompt_database/services/
transaction.rs

1//! Generic transaction wrappers that work directly with [`PgPool`] /
2//! [`PgDbPool`] without going through the dyn-safe trait.
3//!
4//! Copyright (c) systemprompt.io — Business Source License 1.1.
5//! See <https://systemprompt.io> for licensing details.
6
7use crate::error::RepositoryError;
8use crate::repository::PgDbPool;
9use crate::resilience::classify::Outcome;
10use crate::resilience::config::RetryConfig;
11use crate::resilience::retry::retry_async;
12use sqlx::{PgPool, Postgres, Transaction};
13use std::future::Future;
14use std::pin::Pin;
15use std::time::Duration;
16
17pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
18
19pub async fn with_transaction<F, T, E>(pool: &PgPool, f: F) -> Result<T, E>
20where
21    F: for<'c> FnOnce(&'c mut Transaction<'_, Postgres>) -> BoxFuture<'c, Result<T, E>>,
22    E: From<sqlx::Error>,
23{
24    let mut tx = pool.begin().await?;
25    let result = f(&mut tx).await?;
26    tx.commit().await?;
27    Ok(result)
28}
29
30pub async fn with_transaction_retry<F, T>(
31    pool: &PgDbPool,
32    max_retries: u32,
33    f: F,
34) -> Result<T, RepositoryError>
35where
36    T: Send,
37    F: for<'c> Fn(&'c mut Transaction<'_, Postgres>) -> BoxFuture<'c, Result<T, RepositoryError>>
38        + Send
39        + Sync,
40{
41    let cfg = RetryConfig {
42        max_attempts: max_retries.saturating_add(1),
43        base_delay: Duration::from_millis(20),
44        max_delay: Duration::from_millis(640),
45        jitter: false,
46    };
47    let classify = |err: &RepositoryError| {
48        if err.is_serialization_failure() {
49            Outcome::Transient { retry_after: None }
50        } else {
51            Outcome::Permanent
52        }
53    };
54    let attempt = || async {
55        let mut tx = pool.begin().await?;
56        match f(&mut tx).await {
57            Ok(result) => {
58                tx.commit().await?;
59                Ok(result)
60            },
61            Err(e) => {
62                if let Err(rollback_err) = tx.rollback().await {
63                    tracing::error!(error = %rollback_err, "Transaction rollback failed");
64                }
65                Err(e)
66            },
67        }
68    };
69    retry_async(&cfg, "transaction", classify, attempt).await
70}
systemprompt_database/services/transaction.rs

systemprompt_database/services/
transaction.rs
