umbral_core/orm/queryset/
tx.rs

1//! `QuerySetTx` — a QuerySet bound to an open transaction.
2//!
3//! Construction happens in [`super::QuerySet::on_tx`] /
4//! [`super::Manager::on_tx`] using struct-literal syntax against the
5//! `pub(super)` fields. All terminals here mirror their plain-QuerySet
6//! siblings but route their SQL through the borrowed
7//! [`crate::db::Transaction`] so the operations commit or roll back
8//! as a unit with every other operation in the same
9//! `umbral::db::transaction(...)` closure.
10//!
11//! The struct borrows `&mut Transaction` so the borrow checker
12//! enforces that only one `QuerySetTx` uses the transaction at a
13//! time, and that the transaction stays alive for the duration of
14//! each terminal call.
15
16use sea_query::{Expr, Func, PostgresQueryBuilder, SqliteQueryBuilder};
17use sea_query_binder::SqlxBinder;
18
19use crate::orm::{HydrateRelated, Model};
20
21use super::QuerySet;
22use super::errors::GetError;
23use super::write_helpers::{build_insert_one_for, pk_field, serialize_to_map};
24
25/// A `QuerySet` bound to an open transaction. See module docs for
26/// the construction sites and the borrow-checker contract.
27pub struct QuerySetTx<'tx, T> {
28    pub(super) qs: QuerySet<T>,
29    pub(super) tx: &'tx mut crate::db::Transaction,
30}
31
32impl<'tx, T: Model> QuerySetTx<'tx, T> {
33    // -----------------------------------------------------------------------
34    // Read terminals
35    // -----------------------------------------------------------------------
36
37    /// SELECT all matching rows inside the transaction.
38    pub async fn fetch(self) -> Result<Vec<T>, sqlx::Error>
39    where
40        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>
41            + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>
42            + HydrateRelated,
43    {
44        let q = self.qs.build_query_for(self.tx.backend_name());
45        let mut rows = match self.tx.backend_name() {
46            "sqlite" => {
47                let tx = self.tx.as_sqlite_mut().unwrap();
48                let (sql, values) = q.build_sqlx(SqliteQueryBuilder);
49                sqlx::query_as_with::<sqlx::Sqlite, T, _>(&sql, values)
50                    .fetch_all(&mut **tx)
51                    .await?
52            }
53            _ => {
54                let tx = self.tx.as_pg_mut().unwrap();
55                let (sql, values) = q.build_sqlx(PostgresQueryBuilder);
56                sqlx::query_as_with::<sqlx::Postgres, T, _>(&sql, values)
57                    .fetch_all(&mut **tx)
58                    .await?
59            }
60        };
61        // BUG-16 step 2: wire each row's PK into its M2M slots so
62        // junction-table accessors used inside the transaction see
63        // the right parent.
64        for r in &mut rows {
65            r.set_m2m_parent_ids();
66        }
67        Ok(rows)
68    }
69
70    /// SELECT LIMIT 1 and return the first row, if any.
71    pub async fn first(mut self) -> Result<Option<T>, sqlx::Error>
72    where
73        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>
74            + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>
75            + HydrateRelated,
76    {
77        self.qs.query.limit(1);
78        let q = self.qs.build_query_for(self.tx.backend_name());
79        let mut row = match self.tx.backend_name() {
80            "sqlite" => {
81                let tx = self.tx.as_sqlite_mut().unwrap();
82                let (sql, values) = q.build_sqlx(SqliteQueryBuilder);
83                sqlx::query_as_with::<sqlx::Sqlite, T, _>(&sql, values)
84                    .fetch_optional(&mut **tx)
85                    .await?
86            }
87            _ => {
88                let tx = self.tx.as_pg_mut().unwrap();
89                let (sql, values) = q.build_sqlx(PostgresQueryBuilder);
90                sqlx::query_as_with::<sqlx::Postgres, T, _>(&sql, values)
91                    .fetch_optional(&mut **tx)
92                    .await?
93            }
94        };
95        if let Some(r) = row.as_mut() {
96            r.set_m2m_parent_ids();
97        }
98        Ok(row)
99    }
100
101    /// SELECT COUNT(*) inside the transaction.
102    pub async fn count(self) -> Result<i64, sqlx::Error> {
103        let backend = self.tx.backend_name();
104        let mut rebuilt = self.qs.build_query_for(backend);
105        rebuilt.clear_selects();
106        // `sea_query::Asterisk` renders the bare SQL `*` token; `Alias::new("*")`
107        // would render `COUNT("*")` — a quoted identifier Postgres reads as a
108        // column named `*`. Matches the non-transactional count path.
109        rebuilt.expr(Func::count(Expr::col(sea_query::Asterisk)));
110        rebuilt.reset_limit();
111        rebuilt.reset_offset();
112        match backend {
113            "sqlite" => {
114                let tx = self.tx.as_sqlite_mut().unwrap();
115                let (sql, values) = rebuilt.build_sqlx(SqliteQueryBuilder);
116                let (n,): (i64,) = sqlx::query_as_with::<sqlx::Sqlite, (i64,), _>(&sql, values)
117                    .fetch_one(&mut **tx)
118                    .await?;
119                Ok(n)
120            }
121            _ => {
122                let tx = self.tx.as_pg_mut().unwrap();
123                let (sql, values) = rebuilt.build_sqlx(PostgresQueryBuilder);
124                let (n,): (i64,) = sqlx::query_as_with::<sqlx::Postgres, (i64,), _>(&sql, values)
125                    .fetch_one(&mut **tx)
126                    .await?;
127                Ok(n)
128            }
129        }
130    }
131
132    /// Return whether any row matches, inside the transaction.
133    pub async fn exists(mut self) -> Result<bool, sqlx::Error>
134    where
135        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>
136            + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>,
137    {
138        self.qs.query.limit(1);
139        let backend = self.tx.backend_name();
140        let q = self.qs.build_query_for(backend);
141        let row_opt: Option<T> = match backend {
142            "sqlite" => {
143                let tx = self.tx.as_sqlite_mut().unwrap();
144                let (sql, values) = q.build_sqlx(SqliteQueryBuilder);
145                sqlx::query_as_with::<sqlx::Sqlite, T, _>(&sql, values)
146                    .fetch_optional(&mut **tx)
147                    .await?
148            }
149            _ => {
150                let tx = self.tx.as_pg_mut().unwrap();
151                let (sql, values) = q.build_sqlx(PostgresQueryBuilder);
152                sqlx::query_as_with::<sqlx::Postgres, T, _>(&sql, values)
153                    .fetch_optional(&mut **tx)
154                    .await?
155            }
156        };
157        Ok(row_opt.is_some())
158    }
159
160    /// Exactly-one terminal inside the transaction. See [`super::QuerySet::get`].
161    pub async fn get(mut self) -> Result<T, GetError>
162    where
163        T: for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>
164            + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>,
165    {
166        self.qs.query.limit(2);
167        let q = self.qs.build_query_for(self.tx.backend_name());
168        let mut rows: Vec<T> = match self.tx.backend_name() {
169            "sqlite" => {
170                let tx = self.tx.as_sqlite_mut().unwrap();
171                let (sql, values) = q.build_sqlx(SqliteQueryBuilder);
172                sqlx::query_as_with::<sqlx::Sqlite, T, _>(&sql, values)
173                    .fetch_all(&mut **tx)
174                    .await
175                    .map_err(GetError::Sqlx)?
176            }
177            _ => {
178                let tx = self.tx.as_pg_mut().unwrap();
179                let (sql, values) = q.build_sqlx(PostgresQueryBuilder);
180                sqlx::query_as_with::<sqlx::Postgres, T, _>(&sql, values)
181                    .fetch_all(&mut **tx)
182                    .await
183                    .map_err(GetError::Sqlx)?
184            }
185        };
186        match rows.len() {
187            0 => Err(GetError::NotFound),
188            1 => Ok(rows.pop().unwrap()),
189            _ => Err(GetError::MultipleObjectsReturned),
190        }
191    }
192
193    // -----------------------------------------------------------------------
194    // Write terminals
195    // -----------------------------------------------------------------------
196
197    /// DELETE inside the transaction. Returns the number of rows deleted.
198    ///
199    /// On a `#[umbral(soft_delete)]` model this rewrites to
200    /// `UPDATE ... SET deleted_at = NOW()` (plus the on_delete=cascade
201    /// soft-cascade), exactly like the non-transactional `QuerySet::delete`
202    /// — otherwise a `.delete()` that happened to run inside `on_tx()` would
203    /// permanently destroy rows the caller expected to be recoverable.
204    /// `.hard_delete()` opts back into a real DELETE.
205    pub async fn delete(self) -> Result<u64, sqlx::Error> {
206        if self.qs.soft_delete_active && !self.qs.hard_delete {
207            return self.soft_delete_in_tx().await;
208        }
209        let stmt = self.qs.build_delete_for(self.tx.backend_name());
210        match self.tx.backend_name() {
211            "sqlite" => {
212                let tx = self.tx.as_sqlite_mut().unwrap();
213                let (sql, values) = stmt.build_sqlx(SqliteQueryBuilder);
214                let result = sqlx::query_with::<sqlx::Sqlite, _>(&sql, values)
215                    .execute(&mut **tx)
216                    .await?;
217                Ok(result.rows_affected())
218            }
219            _ => {
220                let tx = self.tx.as_pg_mut().unwrap();
221                let (sql, values) = stmt.build_sqlx(PostgresQueryBuilder);
222                let result = sqlx::query_with::<sqlx::Postgres, _>(&sql, values)
223                    .execute(&mut **tx)
224                    .await?;
225                Ok(result.rows_affected())
226            }
227        }
228    }
229
230    /// The soft-delete rewrite of [`Self::delete`], run inside the caller's
231    /// transaction: cascade to any `on_delete = "cascade"` children, then
232    /// stamp `deleted_at = NOW()` on the matched live rows (idempotent —
233    /// never re-stamps an already soft-deleted row). Mirrors
234    /// `QuerySet::soft_delete_update`, minus the private tx it opens.
235    async fn soft_delete_in_tx(self) -> Result<u64, sqlx::Error> {
236        use sea_query::{Alias, Query, Value};
237        let backend = self.tx.backend_name();
238        let now = chrono::Utc::now();
239        let table = crate::db::router::schema_qualified_table(T::TABLE);
240
241        // Cascade first — locate children through the parent's still-live
242        // predicate before the parent is stamped, so no orphaned live child
243        // is left behind.
244        if let Some(pkf) = pk_field::<T>() {
245            let mut sel = Query::select();
246            sel.column(Alias::new(pkf.name)).from(table.clone());
247            for p in &self.qs.predicates {
248                sel.and_where(p.cond_for(backend));
249            }
250            sel.and_where(Expr::col(Alias::new("deleted_at")).is_null());
251            let meta = crate::migrate::ModelMeta::for_::<T>();
252            let mut conn = crate::orm::soft_delete_cascade::CascadeConn::from_tx(self.tx);
253            crate::orm::soft_delete_cascade::cascade_soft_delete(&mut conn, &meta, sel, now)
254                .await?;
255        }
256
257        let mut stmt = Query::update();
258        stmt.table(table);
259        stmt.value(
260            Alias::new("deleted_at"),
261            Value::ChronoDateTimeUtc(Some(Box::new(now))),
262        );
263        for p in &self.qs.predicates {
264            stmt.and_where(p.cond_for(backend));
265        }
266        stmt.and_where(Expr::col(Alias::new("deleted_at")).is_null());
267
268        match backend {
269            "sqlite" => {
270                let tx = self.tx.as_sqlite_mut().unwrap();
271                let (sql, values) = stmt.build_sqlx(SqliteQueryBuilder);
272                let result = sqlx::query_with::<sqlx::Sqlite, _>(&sql, values)
273                    .execute(&mut **tx)
274                    .await?;
275                Ok(result.rows_affected())
276            }
277            _ => {
278                let tx = self.tx.as_pg_mut().unwrap();
279                let (sql, values) = stmt.build_sqlx(PostgresQueryBuilder);
280                let result = sqlx::query_with::<sqlx::Postgres, _>(&sql, values)
281                    .execute(&mut **tx)
282                    .await?;
283                Ok(result.rows_affected())
284            }
285        }
286    }
287
288    /// UPDATE inside the transaction. Takes the same `column → JSON value`
289    /// map as [`super::QuerySet::update_values`].
290    pub async fn update_values(
291        self,
292        values: serde_json::Map<String, serde_json::Value>,
293    ) -> Result<u64, crate::orm::write::WriteError> {
294        let stmt = self.qs.build_update_for(self.tx.backend_name(), &values)?;
295        match self.tx.backend_name() {
296            "sqlite" => {
297                let tx = self.tx.as_sqlite_mut().unwrap();
298                let (sql, values) = stmt.build_sqlx(SqliteQueryBuilder);
299                let result = sqlx::query_with::<sqlx::Sqlite, _>(&sql, values)
300                    .execute(&mut **tx)
301                    .await?;
302                Ok(result.rows_affected())
303            }
304            _ => {
305                let tx = self.tx.as_pg_mut().unwrap();
306                let (sql, values) = stmt.build_sqlx(PostgresQueryBuilder);
307                let result = sqlx::query_with::<sqlx::Postgres, _>(&sql, values)
308                    .execute(&mut **tx)
309                    .await?;
310                Ok(result.rows_affected())
311            }
312        }
313    }
314
315    /// INSERT one row and return the populated row, inside the transaction.
316    ///
317    /// This is the `Manager::create_in_tx` equivalent called through the
318    /// QuerySet API: `Post::objects().on_tx(tx).create(instance).await?`.
319    pub async fn create(self, instance: T) -> Result<T, crate::orm::write::WriteError>
320    where
321        T: serde::Serialize
322            + for<'r> sqlx::FromRow<'r, sqlx::sqlite::SqliteRow>
323            + for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow>
324            + HydrateRelated,
325    {
326        let map = serialize_to_map(&instance)?;
327        let stmt = build_insert_one_for::<T>(self.tx.backend_name(), &map)?;
328        match self.tx.backend_name() {
329            "sqlite" => {
330                let tx = self.tx.as_sqlite_mut().unwrap();
331                let (sql, values) = stmt.build_sqlx(SqliteQueryBuilder);
332                // Classify UNIQUE / FK / NOT NULL / CHECK violations into the
333                // structured `WriteError` variants, symmetric with the non-tx
334                // `QuerySet::create`. Without this a constraint violation inside
335                // a transaction surfaces as an opaque `Sqlx(_)`, so callers that
336                // branch on `WriteError::UniqueViolation` (e.g. the OAuth
337                // username-retry loop) can't tell a collision from a real error.
338                let mut row = sqlx::query_as_with::<sqlx::Sqlite, T, _>(&sql, values)
339                    .fetch_one(&mut **tx)
340                    .await
341                    .map_err(|e| {
342                        crate::orm::validation::classify_sql_error(&e, &map)
343                            .unwrap_or(crate::orm::write::WriteError::Sqlx(e))
344                    })?;
345                row.set_m2m_parent_ids();
346                Ok(row)
347            }
348            _ => {
349                let tx = self.tx.as_pg_mut().unwrap();
350                let (sql, values) = stmt.build_sqlx(PostgresQueryBuilder);
351                let mut row = sqlx::query_as_with::<sqlx::Postgres, T, _>(&sql, values)
352                    .fetch_one(&mut **tx)
353                    .await
354                    .map_err(|e| {
355                        crate::orm::validation::classify_sql_error(&e, &map)
356                            .unwrap_or(crate::orm::write::WriteError::Sqlx(e))
357                    })?;
358                row.set_m2m_parent_ids();
359                Ok(row)
360            }
361        }
362    }
363}
umbral_core/orm/queryset/tx.rs

umbral_core/orm/queryset/
tx.rs
