saddle_db/
transaction.rs

1use std::{future::Future, pin::Pin};
2
3use futures_util::TryStreamExt;
4use saddle_observability::{ActiveCall, CallKind, Observer};
5use sqlx::MySql;
6use tokio::sync::mpsc;
7
8use crate::{
9    CallContext, DbRow, MAX_QUERY_ROWS, MAX_RESULT_BYTES, Result, Statement, WriteResult,
10    cleanup::RollbackJob,
11    error::{
12        map_operation_error, result_limit_exceeded, transaction_commit_failed,
13        transaction_rollback_failed,
14    },
15    row::row_payload_bytes,
16};
17
18pub type TransactionFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T>> + Send + 'a>>;
19
20/// A single-level managed transaction.
21///
22/// There is intentionally no API for beginning another transaction, creating
23/// savepoints, committing, or rolling back. The surrounding `Database` owns
24/// that boundary.
25pub struct Transaction {
26    inner: Option<sqlx::Transaction<'static, MySql>>,
27    observer: Observer,
28    context: CallContext,
29    boundary: Option<ActiveCall>,
30    cleanup: Option<mpsc::UnboundedSender<RollbackJob>>,
31}
32
33impl Transaction {
34    pub(crate) fn new(
35        inner: sqlx::Transaction<'static, MySql>,
36        observer: Observer,
37        boundary: ActiveCall,
38        cleanup: mpsc::UnboundedSender<RollbackJob>,
39    ) -> Self {
40        let context = boundary.context().clone();
41        Self {
42            inner: Some(inner),
43            observer,
44            context,
45            boundary: Some(boundary),
46            cleanup: Some(cleanup),
47        }
48    }
49
50    fn inner(&mut self) -> &mut sqlx::Transaction<'static, MySql> {
51        self.inner
52            .as_mut()
53            .expect("active transaction has an inner connection")
54    }
55
56    pub async fn query_all(&mut self, statement: Statement) -> Result<Vec<DbRow>> {
57        statement.validate()?;
58        let operation = statement.operation().to_owned();
59        let call = self.observer.start_child_call(
60            &self.context,
61            CallKind::Database,
62            "database",
63            "database",
64            operation,
65        );
66        let result = async {
67            let mut stream = statement.query().fetch(&mut **self.inner());
68            let mut rows = Vec::new();
69            let mut result_bytes = 0_usize;
70            while let Some(row) = stream.try_next().await.map_err(map_operation_error)? {
71                if rows.len() == MAX_QUERY_ROWS {
72                    return Err(result_limit_exceeded());
73                }
74                result_bytes = result_bytes.saturating_add(row_payload_bytes(&row)?);
75                if result_bytes > MAX_RESULT_BYTES {
76                    return Err(result_limit_exceeded());
77                }
78                rows.push(DbRow(row));
79            }
80            Ok(rows)
81        }
82        .await;
83        finish_call(call, &result);
84        result
85    }
86
87    pub async fn query_optional(&mut self, statement: Statement) -> Result<Option<DbRow>> {
88        statement.validate()?;
89        let operation = statement.operation().to_owned();
90        let call = self.observer.start_child_call(
91            &self.context,
92            CallKind::Database,
93            "database",
94            "database",
95            operation,
96        );
97        let result = statement
98            .query()
99            .fetch_optional(&mut **self.inner())
100            .await
101            .map_err(map_operation_error)
102            .and_then(|row| {
103                row.map(|row| {
104                    row_payload_bytes(&row)?;
105                    Ok(DbRow(row))
106                })
107                .transpose()
108            });
109        finish_call(call, &result);
110        result
111    }
112
113    pub async fn write(&mut self, statement: Statement) -> Result<WriteResult> {
114        statement.validate()?;
115        let operation = statement.operation().to_owned();
116        let call = self.observer.start_child_call(
117            &self.context,
118            CallKind::Database,
119            "database",
120            "database",
121            operation,
122        );
123        let result = statement
124            .query()
125            .execute(&mut **self.inner())
126            .await
127            .map(|result| WriteResult::new(result.rows_affected(), result.last_insert_id()))
128            .map_err(map_operation_error);
129        finish_call(call, &result);
130        result
131    }
132
133    pub(crate) async fn commit(mut self) -> Result<()> {
134        let inner = self
135            .inner
136            .take()
137            .expect("active transaction has an inner connection");
138        let boundary = self
139            .boundary
140            .take()
141            .expect("active transaction has a boundary");
142        let _cleanup = self.cleanup.take();
143        let commit = self.observer.start_child_call(
144            boundary.context(),
145            CallKind::Transaction,
146            "database",
147            "database",
148            "commit",
149        );
150        match inner.commit().await {
151            Ok(()) => {
152                commit.succeed();
153                boundary.succeed();
154                Ok(())
155            }
156            Err(_) => {
157                let error = transaction_commit_failed();
158                commit.fail(&error);
159                boundary.fail(&error);
160                Err(error)
161            }
162        }
163    }
164
165    pub(crate) async fn rollback(mut self, work_error: crate::SaddleError) -> crate::SaddleError {
166        let inner = self
167            .inner
168            .take()
169            .expect("active transaction has an inner connection");
170        let boundary = self
171            .boundary
172            .take()
173            .expect("active transaction has a boundary");
174        let _cleanup = self.cleanup.take();
175        let rollback = self.observer.start_child_call(
176            boundary.context(),
177            CallKind::Transaction,
178            "database",
179            "database",
180            "rollback",
181        );
182        match inner.rollback().await {
183            Ok(()) => {
184                rollback.succeed();
185                boundary.fail(&work_error);
186                work_error
187            }
188            Err(_) => {
189                let error = transaction_rollback_failed();
190                rollback.fail(&error);
191                boundary.fail(&error);
192                error
193            }
194        }
195    }
196}
197
198impl Drop for Transaction {
199    fn drop(&mut self) {
200        let (Some(inner), Some(boundary), Some(cleanup)) =
201            (self.inner.take(), self.boundary.take(), self.cleanup.take())
202        else {
203            return;
204        };
205        let rollback = self.observer.start_child_call(
206            boundary.context(),
207            CallKind::Transaction,
208            "database",
209            "database",
210            "rollback",
211        );
212        let job = RollbackJob {
213            inner,
214            boundary,
215            rollback,
216        };
217        if let Err(error) = cleanup.send(job) {
218            error.0.fail_without_worker();
219        }
220    }
221}
222
223fn finish_call<T>(call: saddle_observability::ActiveCall, result: &Result<T>) {
224    match result {
225        Ok(_) => call.succeed(),
226        Err(error) => call.fail(error),
227    }
228}
saddle_db/transaction.rs

saddle_db/
transaction.rs
