siderite_orm/queryset/
write.rs

1//! Terminal methods that write: `update`, `delete`, `create`, the
2//! `*_or_create` family and the bulk operations.
3//!
4//! Multi-statement operations run in [`Db::transaction`] (a savepoint when
5//! the queryset's [`Db`] is already a transaction), so they are all-or-nothing.
6
7use super::QuerySet;
8use super::fetch::{is_flat, unordered};
9use crate::Feature;
10use crate::capabilities::RowLocking;
11use crate::db::Db;
12use crate::error::{BackendError, OrmError, QueryError};
13use crate::expr::{Expr, Ident};
14use crate::model::Model;
15use crate::ops::ModelOps;
16use crate::persist::{all_columns, insert_one, insert_values, pk_field, restore_order};
17use crate::plan::SelectExpr;
18use crate::types::DbType;
19use crate::value::Value;
20use crate::write::{DeletePlan, InsertPlan, UpdatePlan, WritePlan};
21
22impl<M: Model> QuerySet<M> {
23    /// `WHERE` clause selecting exactly this queryset's rows.
24    ///
25    /// A plain filter is used directly; anything else (joins, limit, offset,
26    /// distinct) becomes `pk IN (SELECT pk ...)`, since neither `UPDATE` nor
27    /// `DELETE` can carry those clauses.
28    fn write_filter(&self) -> Result<Option<Expr>, QueryError> {
29        let plan = &self.plan;
30        if !plan.grouping.is_empty() || plan.having.is_some() || !plan.compound.is_empty() {
31            return Err(QueryError::InvalidPlan(
32                "cannot update or delete a grouped or combined queryset".into(),
33            ));
34        }
35        if is_flat(plan) && plan.joins.is_empty() {
36            return Ok(plan.filter.clone());
37        }
38        let pk = pk_field::<M>()?.column;
39        let mut keys = unordered(plan);
40        keys.projection = vec![SelectExpr::new(Expr::col(pk), None)];
41        Ok(Some(Expr::col(pk).in_subquery(keys)))
42    }
43
44    /// Set `assignments` (`User::age.set(5)`, `Post::likes.set_expr(Post::likes + 1)`)
45    /// on every matching row; returns the number of rows changed.
46    ///
47    /// Runs one `UPDATE`. `auto_now` columns are not touched: name them in
48    /// `assignments` if they should change.
49    ///
50    /// # Errors
51    /// Capability or backend errors ([`BackendError::Constraint`] for
52    /// violated constraints).
53    pub async fn update(
54        self,
55        assignments: impl IntoIterator<Item = (Ident, Expr)>,
56    ) -> Result<u64, OrmError> {
57        self.ready()?;
58        if self.empty {
59            return Ok(0);
60        }
61        let plan = WritePlan::Update(UpdatePlan {
62            table: M::META.table.into(),
63            assignments: assignments.into_iter().collect(),
64            filter: self.write_filter()?,
65            returning: Vec::new(),
66        });
67        Ok(self.db.execute(&plan).await?.rows_affected)
68    }
69
70    /// Delete every matching row; returns the number of rows removed.
71    ///
72    /// Dependent rows follow the foreign keys' database actions
73    /// (`ON DELETE CASCADE` / `SET NULL` / `RESTRICT`). A `RESTRICT`
74    /// (`OnDelete::Protect`) violation is reported as
75    /// [`BackendError::Constraint`]; nothing is deleted in that case.
76    ///
77    /// # Errors
78    /// Capability or backend errors.
79    pub async fn delete(self) -> Result<u64, OrmError> {
80        self.ready()?;
81        if self.empty {
82            return Ok(0);
83        }
84        let plan = WritePlan::Delete(DeletePlan {
85            table: M::META.table.into(),
86            filter: self.write_filter()?,
87            returning: Vec::new(),
88        });
89        Ok(self.db.execute(&plan).await?.rows_affected)
90    }
91
92    /// Insert `object` and return it as stored (generated key, database
93    /// defaults and stamped timestamps filled in).
94    ///
95    /// # Errors
96    /// Capability or backend errors.
97    pub async fn create(self, mut object: M) -> Result<M, OrmError> {
98        insert_one(&self.db, &mut object).await?;
99        Ok(object)
100    }
101
102    /// The single row matching `predicate`, or a new one built by `default`
103    /// and inserted; the flag says whether it was created.
104    ///
105    /// If a concurrent writer inserts the row first (unique violation), the
106    /// existing row is returned instead. When that re-fetch fails too, the
107    /// [`BackendError::Constraint`] is returned with the re-fetch error
108    /// appended to its message.
109    ///
110    /// # Errors
111    /// [`QueryError::MultipleObjectsReturned`], or capability / backend errors.
112    pub async fn get_or_create(
113        self,
114        predicate: Expr,
115        default: impl FnOnce() -> M + Send,
116    ) -> Result<(M, bool), OrmError> {
117        self.ready()?;
118        self.plan.check(&self.db.capabilities())?;
119        let lookup = self.filter(predicate);
120        match lookup.clone().one().await {
121            Ok(found) => return Ok((found, false)),
122            Err(OrmError::Query(QueryError::DoesNotExist)) => {}
123            Err(other) => return Err(other),
124        }
125        let created = lookup
126            .db
127            .transaction(|tx| async move {
128                let mut object = default();
129                insert_one(&tx, &mut object).await?;
130                Ok::<M, OrmError>(object)
131            })
132            .await;
133        match created {
134            Ok(object) => Ok((object, true)),
135            Err(OrmError::Backend(BackendError::Constraint(message))) => match lookup.one().await {
136                Ok(found) => Ok((found, false)),
137                Err(refetch) => Err(conflict_with_context(message, &refetch)),
138            },
139            Err(other) => Err(other),
140        }
141    }
142
143    /// Apply `update` to the single row matching `predicate` and save it, or
144    /// insert the object built by `create`; the flag says whether it was
145    /// created. Runs in a transaction.
146    ///
147    /// On backends with row locking the lookup is `SELECT ... FOR UPDATE`, so
148    /// concurrent updates of an existing row serialize instead of losing
149    /// writes. When the row is missing, two callers can both try to insert;
150    /// with a unique constraint covering `predicate` the loser's insert fails
151    /// with [`BackendError::Constraint`], and it retries once as lookup and
152    /// update in a new transaction. If that lookup misses too, the constraint
153    /// error is returned with the lookup's error appended.
154    ///
155    /// Backends without row locking (MongoDB) cannot lock the lookup, so two
156    /// concurrent updates of the same existing row can each read the old
157    /// values and the last `save` wins, losing the other's changes. The ORM
158    /// has no conditional (versioned) save; where that matters, express the
159    /// change as an atomic [`update`](Self::update) (e.g.
160    /// `Post::likes.set_expr(Post::likes + 1)`) instead. SQLite serializes
161    /// writers, so it is not affected.
162    ///
163    /// # Errors
164    /// [`QueryError::MultipleObjectsReturned`], or capability / backend errors.
165    pub async fn update_or_create(
166        self,
167        predicate: Expr,
168        create: impl FnOnce() -> M + Send,
169        update: impl FnOnce(&mut M) + Send,
170    ) -> Result<(M, bool), OrmError> {
171        self.ready()?;
172        let lock = self.db.capabilities().row_locking > RowLocking::None;
173        let mut lookup = self.filter(predicate);
174        if lock {
175            lookup = lookup.select_for_update();
176        }
177        // Taken only when the row is found, so after a lost insert race it is
178        // still here for the retry.
179        let mut update = Some(update);
180        let slot = &mut update;
181        let first = lookup
182            .db
183            .clone()
184            .transaction(|tx| {
185                let lookup = lookup.clone();
186                async move {
187                    match lookup.using(&tx).one().await {
188                        Ok(mut found) => {
189                            if let Some(update) = slot.take() {
190                                update(&mut found);
191                            }
192                            found.save(&tx).await?;
193                            Ok((found, false))
194                        }
195                        Err(OrmError::Query(QueryError::DoesNotExist)) => {
196                            let mut object = create();
197                            insert_one(&tx, &mut object).await?;
198                            Ok((object, true))
199                        }
200                        Err(other) => Err(other),
201                    }
202                }
203            })
204            .await;
205        let (Err(OrmError::Backend(BackendError::Constraint(message))), Some(update)) =
206            (&first, update)
207        else {
208            return first;
209        };
210        // The insert lost to a concurrent one: update the row it created.
211        let message = message.clone();
212        lookup
213            .db
214            .clone()
215            .transaction(|tx| async move {
216                let mut found = match lookup.using(&tx).one().await {
217                    Ok(found) => found,
218                    Err(refetch) => return Err(conflict_with_context(message, &refetch)),
219                };
220                update(&mut found);
221                found.save(&tx).await?;
222                Ok((found, false))
223            })
224            .await
225    }
226
227    /// Insert many objects with multi-row `INSERT`s, chunked to the backend's
228    /// parameter limit, in one transaction. Returns the stored rows in input
229    /// order.
230    ///
231    /// Requires `RETURNING`. Rows are matched back to the inputs by key, not
232    /// by the order the database returns them.
233    ///
234    /// # Errors
235    /// Capability or backend errors; nothing is inserted on failure.
236    pub async fn bulk_create(self, objects: Vec<M>) -> Result<Vec<M>, OrmError> {
237        self.ready()?;
238        if objects.is_empty() {
239            return Ok(Vec::new());
240        }
241        let caps = self.db.capabilities();
242        caps.require(Feature::Returning)?;
243        let batches = insert_batches(&objects, caps.max_params);
244        self.db
245            .transaction(|tx| async move {
246                let mut stored = Vec::with_capacity(objects.len());
247                for batch in batches {
248                    let rows = &objects[batch.start..batch.end];
249                    stored.extend(insert_batch(&tx, rows).await?);
250                }
251                Ok::<_, OrmError>(stored)
252            })
253            .await
254    }
255
256    /// Write `columns` of every object in `objects` back to its row, using
257    /// one `UPDATE .. SET col = CASE pk WHEN .. END` per chunk. Returns the
258    /// number of rows changed.
259    ///
260    /// Values are written verbatim (`auto_now` is not refreshed). Objects
261    /// whose row is missing are skipped.
262    ///
263    /// # Errors
264    /// [`QueryError::InvalidPlan`] for an unknown or primary-key column;
265    /// capability or backend errors.
266    pub async fn bulk_update(
267        self,
268        objects: &[M],
269        columns: &[&'static str],
270    ) -> Result<u64, OrmError> {
271        self.ready()?;
272        let pk = pk_field::<M>()?.column;
273        for column in columns {
274            if M::META.column(column).is_none_or(|f| f.primary_key) {
275                return Err(QueryError::InvalidPlan(format!(
276                    "`{column}` is not an updatable column of {}",
277                    M::META.name
278                ))
279                .into());
280            }
281        }
282        if objects.is_empty() || columns.is_empty() {
283            return Ok(0);
284        }
285        let per_chunk = (self.db.capabilities().max_params / (2 * columns.len() + 1)).max(1);
286        self.db
287            .transaction(|tx| async move {
288                let mut changed = 0;
289                for chunk in objects.chunks(per_chunk) {
290                    let plan = bulk_update_plan::<M>(pk, chunk, columns);
291                    changed += tx.execute(&plan).await?.rows_affected;
292                }
293                Ok::<_, OrmError>(changed)
294            })
295            .await
296    }
297}
298
299/// The constraint error of a failed insert, with the error of the follow-up
300/// lookup appended so neither is lost.
301fn conflict_with_context(message: String, refetch: &OrmError) -> OrmError {
302    BackendError::Constraint(format!(
303        "{message} (re-fetch after the conflict failed: {refetch})"
304    ))
305    .into()
306}
307
308/// A run of consecutive objects inserted with one statement.
309struct Batch {
310    start: usize,
311    end: usize,
312}
313
314/// Split `objects` into runs that share a column list (unsaved and saved
315/// objects differ in whether generated columns are written) and fit the
316/// parameter limit.
317fn insert_batches<M: Model>(objects: &[M], max_params: usize) -> Vec<Batch> {
318    let mut batches: Vec<Batch> = Vec::new();
319    let mut previous: Option<Vec<&'static str>> = None;
320    for (i, object) in objects.iter().enumerate() {
321        let columns: Vec<&'static str> = insert_values(object).iter().map(|(c, _)| *c).collect();
322        let rows_per_statement = if columns.is_empty() {
323            1
324        } else {
325            (max_params / columns.len()).max(1)
326        };
327        let extends = previous.as_ref() == Some(&columns)
328            && batches
329                .last()
330                .is_some_and(|b| b.end - b.start < rows_per_statement);
331        match batches.last_mut() {
332            Some(batch) if extends => batch.end = i + 1,
333            _ => batches.push(Batch {
334                start: i,
335                end: i + 1,
336            }),
337        }
338        previous = Some(columns);
339    }
340    batches
341}
342
343async fn insert_batch<M: Model>(db: &Db, objects: &[M]) -> Result<Vec<M>, OrmError> {
344    let rows: Vec<Vec<(&'static str, Value)>> = objects.iter().map(insert_values).collect();
345    let columns: Vec<Ident> = rows
346        .first()
347        .map(|r| r.iter().map(|(c, _)| (*c).into()).collect())
348        .unwrap_or_default();
349    let plan = WritePlan::Insert(InsertPlan {
350        table: M::META.table.into(),
351        columns,
352        rows: rows
353            .into_iter()
354            .map(|r| r.into_iter().map(|(_, v)| v).collect())
355            .collect(),
356        returning: all_columns::<M>(),
357    });
358    let result = db.execute(&plan).await?;
359    if result.returning.len() != objects.len() {
360        return Err(QueryError::Model(format!(
361            "inserted {} rows but the database returned {}",
362            objects.len(),
363            result.returning.len()
364        ))
365        .into());
366    }
367    let stored = result
368        .returning
369        .iter()
370        .map(|row| M::from_row(row, ""))
371        .collect::<Result<Vec<_>, _>>()?;
372    Ok(restore_order(objects, stored))
373}
374
375fn bulk_update_plan<M: Model>(
376    pk: &'static str,
377    chunk: &[M],
378    columns: &[&'static str],
379) -> WritePlan {
380    let assignments = columns
381        .iter()
382        .map(|column| {
383            let branches = chunk
384                .iter()
385                .filter_map(|object| {
386                    let value = object.to_values().into_iter().find(|(c, _)| c == column)?.1;
387                    let matches_row = Expr::col(pk).eq(Expr::Value(object.pk().to_value()));
388                    Some((matches_row, Expr::Value(value)))
389                })
390                .collect();
391            let case = Expr::Case {
392                branches,
393                otherwise: Some(Box::new(Expr::col(*column))),
394            };
395            ((*column).into(), case)
396        })
397        .collect();
398    WritePlan::Update(UpdatePlan {
399        table: M::META.table.into(),
400        assignments,
401        filter: Some(Expr::col(pk).is_in(chunk.iter().map(|o| o.pk().to_value()))),
402        returning: Vec::new(),
403    })
404}
siderite_orm/queryset/write.rs

siderite_orm/queryset/
write.rs
