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: &PgDbPool, 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_raw<F, T, E>(pool: &PgPool, f: F) -> Result<T, E>
31where
32    F: for<'c> FnOnce(&'c mut Transaction<'_, Postgres>) -> BoxFuture<'c, Result<T, E>>,
33    E: From<sqlx::Error>,
34{
35    let mut tx = pool.begin().await?;
36    let result = f(&mut tx).await?;
37    tx.commit().await?;
38    Ok(result)
39}
40
41pub async fn with_transaction_retry<F, T>(
42    pool: &PgDbPool,
43    max_retries: u32,
44    f: F,
45) -> Result<T, RepositoryError>
46where
47    T: Send,
48    F: for<'c> Fn(&'c mut Transaction<'_, Postgres>) -> BoxFuture<'c, Result<T, RepositoryError>>
49        + Send
50        + Sync,
51{
52    let cfg = RetryConfig {
53        max_attempts: max_retries.saturating_add(1),
54        base_delay: Duration::from_millis(20),
55        max_delay: Duration::from_millis(640),
56        jitter: false,
57    };
58    let classify = |err: &RepositoryError| {
59        if is_retriable_error(err) {
60            Outcome::Transient { retry_after: None }
61        } else {
62            Outcome::Permanent
63        }
64    };
65    let attempt = || async {
66        let mut tx = pool.begin().await?;
67        match f(&mut tx).await {
68            Ok(result) => {
69                tx.commit().await?;
70                Ok(result)
71            },
72            Err(e) => {
73                if let Err(rollback_err) = tx.rollback().await {
74                    tracing::error!(error = %rollback_err, "Transaction rollback failed");
75                }
76                Err(e)
77            },
78        }
79    };
80    retry_async(&cfg, "transaction", classify, attempt).await
81}
82
83fn is_retriable_error(error: &RepositoryError) -> bool {
84    match error {
85        RepositoryError::Database(sqlx_error) => {
86            sqlx_error.as_database_error().is_some_and(|db_error| {
87                let code = db_error.code().map(|c| c.to_string());
88                matches!(code.as_deref(), Some("40001" | "40P01"))
89            })
90        },
91        RepositoryError::NotFound(_)
92        | RepositoryError::Constraint(_)
93        | RepositoryError::Serialization(_)
94        | RepositoryError::InvalidArgument(_)
95        | RepositoryError::InvalidState(_)
96        | RepositoryError::Internal(_)
97        | RepositoryError::QueryExecution(_) => false,
98    }
99}
systemprompt_database/services/transaction.rs

systemprompt_database/services/
transaction.rs
