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 {}