editoast: change Model traits CreateBatch,WithKey to use editoast_mod…

…els::Error Signed-off-by: Leo Valais <[email protected]>
OpenRailAssociation · Jan 28, 2025 · 2b56d3e · 2b56d3e
1 parent 57b6bbe
commit 2b56d3e
Show file tree

Hide file tree

Showing 7 changed files with 94 additions and 73 deletions.
diff --git a/editoast/editoast_derive/src/model/codegen/create_batch_impl.rs b/editoast/editoast_derive/src/model/codegen/create_batch_impl.rs
@@ -42,23 +42,25 @@ impl ToTokens for CreateBatchImpl {
                     .returning((#(dsl::#columns,)*))
                     .load_stream::<#row>(conn.write().await.deref_mut())
                     .await
-                    .map(|s| s.map_ok(<#model as Model>::from_row).try_collect::<Vec<_>>())?
-                    .await?
+                    .map_err(|e| <#model as crate::models::Model>::Error::from(editoast_models::model::Error::from(e)))?
+                    .map_ok(<#model as Model>::from_row)
+                    .try_collect::<Vec<_>>()
+                    .await
+                    .map_err(|e| <#model as crate::models::Model>::Error::from(editoast_models::model::Error::from(e)))?
             },
         };
 
         tokens.extend(quote! {
             #[automatically_derived]
-            #[async_trait::async_trait]
-            impl crate::models::CreateBatch<#changeset> for #model {
+            impl crate::models::CreateBatch for #model {
                 #[tracing::instrument(name = #span_name, skip_all, err)]
                 async fn create_batch<
-                    I: std::iter::IntoIterator<Item = #changeset> + Send + 'async_trait,
+                    I: std::iter::IntoIterator<Item = #changeset> + Send,
                     C: Default + std::iter::Extend<Self> + Send + std::fmt::Debug,
                 >(
                     conn: &mut editoast_models::DbConnection,
                     values: I,
-                ) -> crate::error::Result<C> {
+                ) -> std::result::Result<C, <#model as crate::models::Model>::Error> {
                     use crate::models::Model;
                     use #table_mod::dsl;
                     use std::ops::DerefMut;

diff --git a/editoast/editoast_derive/src/model/codegen/create_batch_with_key_impl.rs b/editoast/editoast_derive/src/model/codegen/create_batch_with_key_impl.rs
@@ -47,29 +47,28 @@ impl ToTokens for CreateBatchWithKeyImpl {
                     .returning((#(dsl::#columns,)*))
                     .load_stream::<#row>(conn.write().await.deref_mut())
                     .await
-                    .map(|s| {
-                        s.map_ok(|row| {
-                            let model = <#model as Model>::from_row(row);
-                            (model.get_id(), model)
-                        })
-                        .try_collect::<Vec<_>>()
-                    })?
-                    .await?
+                    .map_err(|e| <#model as crate::models::Model>::Error::from(editoast_models::model::Error::from(e)))?
+                    .map_ok(|row| {
+                        let model = <#model as Model>::from_row(row);
+                        (model.get_id(), model)
+                    })
+                    .try_collect::<Vec<_>>()
+                    .await
+                    .map_err(|e| <#model as crate::models::Model>::Error::from(editoast_models::model::Error::from(e)))?
             },
         };
 
         tokens.extend(quote! {
             #[automatically_derived]
-            #[async_trait::async_trait]
-            impl crate::models::CreateBatchWithKey<#changeset, #ty> for #model {
+            impl crate::models::CreateBatchWithKey<#ty> for #model {
                 #[tracing::instrument(name = #span_name, skip_all, err)]
                 async fn create_batch_with_key<
-                    I: std::iter::IntoIterator<Item = #changeset> + Send + 'async_trait,
+                    I: std::iter::IntoIterator<Item = #changeset> + Send,
                     C: Default + std::iter::Extend<(#ty, Self)> + Send + std::fmt::Debug,
                 >(
                     conn: &mut editoast_models::DbConnection,
                     values: I,
-                ) -> crate::error::Result<C> {
+                ) -> std::result::Result<C, <#model as crate::models::Model>::Error> {
                     use crate::models::Identifiable;
                     use crate::models::Model;
                     use std::ops::DerefMut;

diff --git a/editoast/editoast_derive/src/snapshots/editoast_derive__model__tests__construction.snap b/editoast/editoast_derive/src/snapshots/editoast_derive__model__tests__construction.snap
@@ -611,13 +611,15 @@ impl crate::models::prelude::Count for Document {
     }
 }
 #[automatically_derived]
-#[async_trait::async_trait]
-impl crate::models::CreateBatch<DocumentChangeset> for Document {
+impl crate::models::CreateBatch for Document {
     #[tracing::instrument(name = "model:create_batch<Document>", skip_all, err)]
     async fn create_batch<
-        I: std::iter::IntoIterator<Item = DocumentChangeset> + Send + 'async_trait,
+        I: std::iter::IntoIterator<Item = DocumentChangeset> + Send,
         C: Default + std::iter::Extend<Self> + Send + std::fmt::Debug,
-    >(conn: &mut editoast_models::DbConnection, values: I) -> crate::error::Result<C> {
+    >(
+        conn: &mut editoast_models::DbConnection,
+        values: I,
+    ) -> std::result::Result<C, <Document as crate::models::Model>::Error> {
         use crate::models::Model;
         use editoast_models::tables::osrd_infra_document::dsl;
         use std::ops::DerefMut;
@@ -638,12 +640,15 @@ impl crate::models::CreateBatch<DocumentChangeset> for Document {
                         .returning((dsl::id, dsl::content_type, dsl::data))
                         .load_stream::<DocumentRow>(conn.write().await.deref_mut())
                         .await
-                        .map(|s| {
-                            s
-                                .map_ok(<Document as Model>::from_row)
-                                .try_collect::<Vec<_>>()
-                        })?
-                        .await?
+                        .map_err(|e| <Document as crate::models::Model>::Error::from(
+                            editoast_models::model::Error::from(e),
+                        ))?
+                        .map_ok(<Document as Model>::from_row)
+                        .try_collect::<Vec<_>>()
+                        .await
+                        .map_err(|e| <Document as crate::models::Model>::Error::from(
+                            editoast_models::model::Error::from(e),
+                        ))?
                 };
                 result.extend(chunk_result);
             }
@@ -652,13 +657,15 @@ impl crate::models::CreateBatch<DocumentChangeset> for Document {
     }
 }
 #[automatically_derived]
-#[async_trait::async_trait]
-impl crate::models::CreateBatchWithKey<DocumentChangeset, (String)> for Document {
+impl crate::models::CreateBatchWithKey<(String)> for Document {
     #[tracing::instrument(name = "model:create_batch_with_key<Document>", skip_all, err)]
     async fn create_batch_with_key<
-        I: std::iter::IntoIterator<Item = DocumentChangeset> + Send + 'async_trait,
+        I: std::iter::IntoIterator<Item = DocumentChangeset> + Send,
         C: Default + std::iter::Extend<((String), Self)> + Send + std::fmt::Debug,
-    >(conn: &mut editoast_models::DbConnection, values: I) -> crate::error::Result<C> {
+    >(
+        conn: &mut editoast_models::DbConnection,
+        values: I,
+    ) -> std::result::Result<C, <Document as crate::models::Model>::Error> {
         use crate::models::Identifiable;
         use crate::models::Model;
         use std::ops::DerefMut;
@@ -680,14 +687,18 @@ impl crate::models::CreateBatchWithKey<DocumentChangeset, (String)> for Document
                         .returning((dsl::id, dsl::content_type, dsl::data))
                         .load_stream::<DocumentRow>(conn.write().await.deref_mut())
                         .await
-                        .map(|s| {
-                            s.map_ok(|row| {
-                                    let model = <Document as Model>::from_row(row);
-                                    (model.get_id(), model)
-                                })
-                                .try_collect::<Vec<_>>()
-                        })?
-                        .await?
+                        .map_err(|e| <Document as crate::models::Model>::Error::from(
+                            editoast_models::model::Error::from(e),
+                        ))?
+                        .map_ok(|row| {
+                            let model = <Document as Model>::from_row(row);
+                            (model.get_id(), model)
+                        })
+                        .try_collect::<Vec<_>>()
+                        .await
+                        .map_err(|e| <Document as crate::models::Model>::Error::from(
+                            editoast_models::model::Error::from(e),
+                        ))?
                 };
                 result.extend(chunk_result);
             }
@@ -696,13 +707,15 @@ impl crate::models::CreateBatchWithKey<DocumentChangeset, (String)> for Document
     }
 }
 #[automatically_derived]
-#[async_trait::async_trait]
-impl crate::models::CreateBatchWithKey<DocumentChangeset, (i64)> for Document {
+impl crate::models::CreateBatchWithKey<(i64)> for Document {
     #[tracing::instrument(name = "model:create_batch_with_key<Document>", skip_all, err)]
     async fn create_batch_with_key<
-        I: std::iter::IntoIterator<Item = DocumentChangeset> + Send + 'async_trait,
+        I: std::iter::IntoIterator<Item = DocumentChangeset> + Send,
         C: Default + std::iter::Extend<((i64), Self)> + Send + std::fmt::Debug,
-    >(conn: &mut editoast_models::DbConnection, values: I) -> crate::error::Result<C> {
+    >(
+        conn: &mut editoast_models::DbConnection,
+        values: I,
+    ) -> std::result::Result<C, <Document as crate::models::Model>::Error> {
         use crate::models::Identifiable;
         use crate::models::Model;
         use std::ops::DerefMut;
@@ -724,14 +737,18 @@ impl crate::models::CreateBatchWithKey<DocumentChangeset, (i64)> for Document {
                         .returning((dsl::id, dsl::content_type, dsl::data))
                         .load_stream::<DocumentRow>(conn.write().await.deref_mut())
                         .await
-                        .map(|s| {
-                            s.map_ok(|row| {
-                                    let model = <Document as Model>::from_row(row);
-                                    (model.get_id(), model)
-                                })
-                                .try_collect::<Vec<_>>()
-                        })?
-                        .await?
+                        .map_err(|e| <Document as crate::models::Model>::Error::from(
+                            editoast_models::model::Error::from(e),
+                        ))?
+                        .map_ok(|row| {
+                            let model = <Document as Model>::from_row(row);
+                            (model.get_id(), model)
+                        })
+                        .try_collect::<Vec<_>>()
+                        .await
+                        .map_err(|e| <Document as crate::models::Model>::Error::from(
+                            editoast_models::model::Error::from(e),
+                        ))?
                 };
                 result.extend(chunk_result);
             }

diff --git a/editoast/src/models/mod.rs b/editoast/src/models/mod.rs
@@ -61,14 +61,18 @@ mod tests {
     use super::prelude::*;
 
     #[derive(Debug, Default, Clone, Model, PartialEq, Eq)]
-    #[model(table = editoast_models::tables::document)]
+    #[model(table = editoast_models::tables::document, error = DocError)]
     #[model(gen(ops = crud, batch_ops = crud, list))]
     struct Document {
         id: i64,
         content_type: String,
         data: Vec<u8>,
     }
 
+    #[derive(Debug, thiserror::Error)]
+    #[error("oh no {0}")]
+    struct DocError(#[from] editoast_models::model::Error);
+
     #[rstest]
     async fn test_batch() {
         let pool = DbConnectionPoolV2::for_tests();

diff --git a/editoast/src/models/prelude/create.rs b/editoast/src/models/prelude/create.rs
@@ -1,14 +1,14 @@
 use std::fmt::Debug;
+use std::result::Result;
 
 use editoast_models::model;
 use editoast_models::DbConnection;
 
 use crate::error::EditoastError;
-use crate::error::Result;
 
 use super::Model;
 
-/// Describes how a [Model](super::Model) can be created in the database
+/// Describes how a [Model] can be created in the database
 ///
 /// You can implement this type manually but its recommended to use the `Model`
 /// derive macro instead.
@@ -30,15 +30,11 @@ pub trait Create<M: Model>: Sized {
     }
 }
 
-/// Describes how a [Model](super::Model) can be created in the database given a batch of its changesets
+/// Describes how a [Model] can be created in the database given a batch of its changesets
 ///
 /// You can implement this type manually but its recommended to use the `Model`
 /// derive macro instead.
-#[async_trait::async_trait]
-pub trait CreateBatch<Cs>: Sized
-where
-    Cs: Send,
-{
+pub trait CreateBatch: Model {
     /// Creates a batch of rows in the database given an iterator of changesets
     ///
     /// Returns a collection of the created rows.
@@ -52,34 +48,32 @@ where
     /// assert_eq!(docs.len(), 5);
     /// ```
     async fn create_batch<
-        I: IntoIterator<Item = Cs> + Send + 'async_trait,
+        I: IntoIterator<Item = Self::Changeset> + Send,
         C: Default + std::iter::Extend<Self> + Send + Debug,
     >(
         conn: &mut DbConnection,
         values: I,
-    ) -> Result<C>;
+    ) -> Result<C, Self::Error>;
 }
 
-/// Describes how a [Model](super::Model) can be created in the database given a batch of its changesets
+/// Describes how a [Model] can be created in the database given a batch of its changesets
 ///
 /// This trait is similar to [CreateBatch] but the returned models are paired with their key.
 /// There is two different traits because Rust cannot infer the type of the key when using
 /// [CreateBatch::create_batch] with a model that has more than one identifier.
 ///
 /// You can implement this type manually but its recommended to use the `Model`
 /// derive macro instead.
-#[async_trait::async_trait]
-pub trait CreateBatchWithKey<Cs, K>: Sized
+pub trait CreateBatchWithKey<K>: Model
 where
-    Cs: Send,
     K: Send + Clone,
 {
     /// Just like [CreateBatch::create_batch] but the returned models are paired with their key
     async fn create_batch_with_key<
-        I: IntoIterator<Item = Cs> + Send + 'async_trait,
+        I: IntoIterator<Item = Self::Changeset> + Send,
         C: Default + std::iter::Extend<(K, Self)> + Send + Debug,
     >(
         conn: &mut DbConnection,
         values: I,
-    ) -> Result<C>;
+    ) -> Result<C, Self::Error>;
 }
diff --git a/editoast/src/views/temporary_speed_limits.rs b/editoast/src/views/temporary_speed_limits.rs
@@ -173,7 +173,9 @@ async fn create_temporary_speed_limit_group(
         .into_iter()
         .map(|speed_limit| speed_limit.into_temporary_speed_limit_changeset(group_id))
         .collect::<Vec<_>>();
-    let _: Vec<_> = TemporarySpeedLimit::create_batch(conn, speed_limits_changesets).await?;
+    let _: Vec<_> = TemporarySpeedLimit::create_batch(conn, speed_limits_changesets)
+        .await
+        .map_err(TemporarySpeedLimitError::from)?;
 
     Ok(Json(TemporarySpeedLimitCreateResponse { group_id }))
 }

diff --git a/editoast/src/views/work_schedules.rs b/editoast/src/views/work_schedules.rs
@@ -202,8 +202,9 @@ async fn create(
             work_schedule.into_work_schedule_changeset(work_schedule_group.work_schedule_group_id)
         })
         .collect::<Vec<_>>();
-    let _work_schedules: Vec<_> =
-        WorkSchedule::create_batch(conn, work_schedules_changesets).await?;
+    let _work_schedules: Vec<_> = WorkSchedule::create_batch(conn, work_schedules_changesets)
+        .await
+        .map_err(WorkScheduleError::from)?;
 
     Ok(Json(WorkScheduleCreateResponse {
         work_schedule_group_id: work_schedule_group.work_schedule_group_id,
@@ -454,7 +455,9 @@ async fn put_in_group(
                 .map(|work_schedule| work_schedule.into_work_schedule_changeset(group_id))
                 .collect::<Vec<_>>();
             let work_schedules =
-                WorkSchedule::create_batch(&mut conn.clone(), work_schedules_changesets).await?;
+                WorkSchedule::create_batch(&mut conn.clone(), work_schedules_changesets)
+                    .await
+                    .map_err(WorkScheduleError::from)?;
 
             Ok(Json(work_schedules))
         })