sqlx_utils/traits/repository/
save.rs

1//! Trait for automatically insert or update based on the presence of an ID.
2
3use crate::prelude::Database;
4use crate::traits::{InsertableRepository, Model, UpdatableRepository};
5use crate::utils::{BatchOperator, DEFAULT_BATCH_SIZE};
6use sqlx::Executor;
7
8/// Trait for repositories that can intelligently save records by either inserting or updating them.
9///
10/// The `SaveRepository` trait combines [`InsertableRepository`] and [`UpdatableRepository`]
11/// to provide a higher-level interface for persisting models. It automatically determines
12/// whether to insert or update a record based on whether the model has an ID, allowing
13/// for simpler code when the operation type doesn't matter to the caller.
14///
15/// # Type Parameters
16///
17/// * `M` - The model type that this repository saves. Must implement the [`Model`] trait.
18///
19/// # Examples
20///
21/// The trait is automatically implemented for any repository that implements both
22/// [`InsertableRepository`] and [`UpdatableRepository`]:
23///
24/// ```rust
25/// # use sqlx_utils::traits::{Model, Repository, InsertableRepository, UpdatableRepository, SaveRepository};
26/// # use sqlx_utils::types::{Pool, Query};
27/// # struct User { id: Option<i32>, name: String }
28/// # impl Model for User {
29/// #     type Id = i32;
30/// #     fn get_id(&self) -> Option<Self::Id> { self.id }
31/// # }
32/// # struct UserRepository { pool: Pool }
33/// # impl Repository<User> for UserRepository {
34/// #     fn pool(&self) -> &Pool { &self.pool }
35/// # }
36///
37/// impl InsertableRepository<User> for UserRepository {
38///     fn insert_query(user: &User) -> Query<'_> {
39///         sqlx::query("INSERT INTO users (name) VALUES ($1)")
40///             .bind(&user.name)
41///     }
42/// }
43///
44/// impl UpdatableRepository<User> for UserRepository {
45///     fn update_query(user: &User) -> Query<'_> {
46///         sqlx::query("UPDATE users SET name = $1 WHERE id = $2")
47///             .bind(&user.name)
48///             .bind(user.id.unwrap())
49///     }
50/// }
51///
52/// // SaveRepository is automatically implemented
53///
54/// // Usage
55/// # async fn example(repo: &UserRepository) -> sqlx_utils::Result<()> {
56/// // Create a new user (no ID)
57/// let new_user = User { id: None, name: String::from("Alice") };
58/// repo.save(new_user).await?; // Will insert
59///
60/// // Update an existing user (has ID)
61/// let existing_user = User { id: Some(1), name: String::from("Updated Alice") };
62/// repo.save_ref(&existing_user).await?; // Will update
63///
64/// // Save a mixed batch of new and existing users
65/// let users = vec![
66///     User { id: None, name: String::from("Bob") },        // Will insert
67///     User { id: Some(2), name: String::from("Charlie") }  // Will update
68/// ];
69/// repo.save_all(users).await?; // Automatically sorts and batches operations
70/// # Ok(())
71/// # }
72/// ```
73///
74/// # Implementation Notes
75///
76/// 1. This trait is automatically implemented for any type that implements both
77///    [`InsertableRepository`] and [`UpdatableRepository`]
78/// 2. The [`save`](SaveRepository::save_with_executor) method checks [`Model::get_id()`] to determine whether to insert or update
79/// 3. The batch methods intelligently sort models into separate insert and update operations
80/// 4. Where possible, insert and update operations within a batch are executed concurrently
81///    for optimal performance
82#[diagnostic::on_unimplemented(
83    note = "Type `{Self}` does not implement the `SaveRepository<{M}>` trait",
84    label = "this type cannot automatically save `{M}` records",
85    message = "`{Self}` must implement both `InsertableRepository<{M}>` and `UpdatableRepository<{M}>` to gain `SaveRepository<{M}>` capabilities"
86)]
87#[async_trait::async_trait]
88pub trait SaveRepository<M: Model>: InsertableRepository<M> + UpdatableRepository<M> {
89    /// Intelligently persists a model instance by either inserting or updating using the [`Executor`] `tx`.
90    ///
91    /// This method determines the appropriate operation based on whether the model
92    /// has an ID:
93    /// - If the model has no ID, it performs an insertion
94    /// - If the model has an ID, it performs an update
95    ///
96    /// # Parameters
97    ///
98    /// * `tx` - The executor to use for the query
99    /// * `model` - A reference to the model instance to save
100    ///
101    /// # Returns
102    ///
103    /// * [`crate::Result<()>`](crate::Result) - Success if the operation was executed, or an error if it failed
104    ///
105    /// # Example
106    ///
107    /// ```ignore
108    /// async fn save_user(repo: &UserRepository, user: &User) -> crate::Result<()> {
109    ///     repo.save(user).await // Will insert or update based on user.id
110    /// }
111    /// ```
112    #[inline]
113    #[cfg_attr(feature = "log_err", tracing::instrument(skip_all, level = "debug", parent = &(Self::repository_span()), name = "save", err))]
114    #[cfg_attr(not(feature = "log_err"), tracing::instrument(skip_all, level = "debug", parent = &(Self::repository_span()), name = "save"))]
115    async fn save_with_executor<'c, E>(&self, tx: E, model: M) -> crate::Result<M>
116    where
117        M: 'async_trait,
118        E: Executor<'c, Database = Database> + Send,
119    {
120        if model.get_id().is_none() {
121            <Self as InsertableRepository<M>>::insert_with_executor(self, tx, model).await
122        } else {
123            <Self as UpdatableRepository<M>>::update_with_executor(self, tx, model).await
124        }
125    }
126
127    /// Intelligently persists a model instance by either inserting or updating using the [`Executor`] `tx`.
128    ///
129    /// This method determines the appropriate operation based on whether the model
130    /// has an ID:
131    /// - If the model has no ID, it performs an insertion
132    /// - If the model has an ID, it performs an update
133    ///
134    /// # Parameters
135    ///
136    /// * `tx` - The executor to use for the query
137    /// * `model` - A reference to the model instance to save
138    ///
139    /// # Returns
140    ///
141    /// * [`crate::Result<()>`](crate::Result) - Success if the operation was executed, or an error if it failed
142    ///
143    /// # Example
144    ///
145    /// ```ignore
146    /// async fn save_user(repo: &UserRepository, user: &User) -> crate::Result<()> {
147    ///     repo.save(user).await // Will insert or update based on user.id
148    /// }
149    /// ```
150    #[inline]
151    #[cfg_attr(feature = "log_err", tracing::instrument(skip_all, level = "debug", parent = &(Self::repository_span()), name = "save", err))]
152    #[cfg_attr(not(feature = "log_err"), tracing::instrument(skip_all, level = "debug", parent = &(Self::repository_span()), name = "save"))]
153    async fn save_ref_with_executor<'c, E>(&self, tx: E, model: &M) -> crate::Result<()>
154    where
155        M: 'async_trait,
156        E: Executor<'c, Database = Database> + Send,
157    {
158        if model.get_id().is_none() {
159            <Self as InsertableRepository<M>>::insert_ref_with_executor(self, tx, model).await
160        } else {
161            <Self as UpdatableRepository<M>>::update_ref_with_executor(self, tx, model).await
162        }
163    }
164
165    /// Intelligently persists a model instance by either inserting or updating.
166    ///
167    /// This method determines the appropriate operation based on whether the model
168    /// has an ID:
169    /// - If the model has no ID, it performs an insertion
170    /// - If the model has an ID, it performs an update
171    ///
172    /// # Parameters
173    ///
174    /// * `model` - A reference to the model instance to save
175    ///
176    /// # Returns
177    ///
178    /// * [`crate::Result<()>`](crate::Result) - Success if the operation was executed, or an error if it failed
179    ///
180    /// # Example
181    ///
182    /// ```ignore
183    /// async fn save_user(repo: &UserRepository, user: &User) -> crate::Result<()> {
184    ///     repo.save(user).await // Will insert or update based on user.id
185    /// }
186    /// ```
187    #[inline(always)]
188    async fn save(&self, model: M) -> crate::Result<M>
189    where
190        M: 'async_trait,
191    {
192        self.save_with_executor(self.pool(), model).await
193    }
194
195    /// Intelligently persists a model instance by either inserting or updating.
196    ///
197    /// This method determines the appropriate operation based on whether the model
198    /// has an ID:
199    /// - If the model has no ID, it performs an insertion
200    /// - If the model has an ID, it performs an update
201    ///
202    /// # Parameters
203    ///
204    /// * `model` - A reference to the model instance to save
205    ///
206    /// # Returns
207    ///
208    /// * [`crate::Result<()>`](crate::Result) - Success if the operation was executed, or an error if it failed
209    ///
210    /// # Example
211    ///
212    /// ```ignore
213    /// async fn save_user(repo: &UserRepository, user: &User) -> crate::Result<()> {
214    ///     repo.save(user).await // Will insert or update based on user.id
215    /// }
216    /// ```
217    #[inline(always)]
218    async fn save_ref(&self, model: &M) -> crate::Result<()>
219    where
220        M: 'async_trait,
221    {
222        self.save_ref_with_executor(self.pool(), model).await
223    }
224
225    /// Saves multiple models using the default batch size.
226    ///
227    /// This is a convenience wrapper around [`save_batch`](Self::save_batch) that uses [`DEFAULT_BATCH_SIZE`].
228    /// It provides a simpler interface for bulk save operations when the default batch
229    /// size is appropriate.
230    ///
231    /// # Parameters
232    ///
233    /// * `models` - An iterator yielding model instances to save
234    ///
235    /// # Returns
236    ///
237    /// * [`crate::Result<()>`](crate::Result) - Success if all operations were executed, or an error if any failed
238    #[inline]
239    async fn save_all<I>(&self, models: I) -> crate::Result<()>
240    where
241        I: IntoIterator<Item = M> + Send + 'async_trait,
242        I::IntoIter: Send,
243    {
244        <Self as SaveRepository<M>>::save_batch::<DEFAULT_BATCH_SIZE, I>(self, models).await
245    }
246
247    /// Performs an intelligent batched save operation with a specified batch size.
248    ///
249    /// This is the most sophisticated batch operation, efficiently handling both
250    /// insertions and updates in the same operation. It sorts models based on
251    /// whether they need insertion or update, then processes them optimally.
252    ///
253    /// # Type Parameters
254    ///
255    /// * `N` - The size of each batch to process
256    ///
257    /// # Parameters
258    ///
259    /// * `models` - An iterator yielding model instances to save
260    ///
261    /// # Returns
262    ///
263    /// * [`crate::Result<()>`](crate::Result) - Success if all batches were processed, or an error if any operation failed
264    ///
265    /// # Implementation Details
266    ///
267    /// The method:
268    /// 1. Splits each batch into models requiring insertion vs update
269    /// 2. Processes insertions and updates concurrently when possible
270    /// 3. Handles empty cases efficiently
271    /// 4. Maintains transactional integrity within each batch
272    ///
273    /// # Performance Features
274    ///
275    /// - Concurrent processing of inserts and updates
276    /// - Efficient batch size management
277    /// - Smart handling of empty cases
278    /// - Transaction management for data consistency
279    ///
280    /// # Performance Considerations
281    ///
282    /// Consider batch size carefully:
283    /// - Too small: More overhead from multiple transactions
284    /// - Too large: Higher memory usage and longer transactions times
285    #[cfg_attr(feature = "log_err", tracing::instrument(skip_all, level = "debug", parent = &(Self::repository_span()), name = "save_batch", err))]
286    #[cfg_attr(not(feature = "log_err"), tracing::instrument(skip_all, level = "debug", parent = &(Self::repository_span()), name = "save_batch"))]
287    async fn save_batch<const N: usize, I>(&self, models: I) -> crate::Result<()>
288    where
289        I: IntoIterator<Item = M> + Send + 'async_trait,
290        I::IntoIter: Send,
291        M: 'async_trait,
292    {
293        let span = tracing::Span::current();
294        span.record("BATCH_SIZE", N);
295
296        let op = BatchOperator::<M, N>::execute_batch(models, |batch| async {
297            let mut update = Vec::new();
298            let mut insert = Vec::new();
299
300            for model in batch {
301                if model.get_id().is_some() {
302                    update.push(model);
303                } else {
304                    insert.push(model);
305                }
306            }
307
308            match (update.is_empty(), insert.is_empty()) {
309                (false, false) => {
310                    futures::try_join!(
311                        <Self as UpdatableRepository<M>>::update_batch::<N, Vec<M>>(self, update),
312                        <Self as InsertableRepository<M>>::insert_batch::<N, Vec<M>>(self, insert)
313                    )?;
314                }
315                (false, true) => {
316                    <Self as UpdatableRepository<M>>::update_batch::<N, Vec<M>>(self, update)
317                        .await?;
318                }
319                (true, false) => {
320                    <Self as InsertableRepository<M>>::insert_batch::<N, Vec<M>>(self, insert)
321                        .await?;
322                }
323                (true, true) => {}
324            }
325
326            Ok(())
327        });
328
329        op.await
330    }
331}
332
333#[async_trait::async_trait]
334impl<M: Model, T: InsertableRepository<M> + UpdatableRepository<M>> SaveRepository<M> for T {}