editoast: add models for timetable and train schedule v2

Co-authored-by: Youness Chrifi Alaoui <[email protected]>

Author: flomonster

editoast/src/modelsv2/mod.rs

@@ -2,11 +2,12 @@ pub mod documents;
 pub mod infra_objects;
 pub mod projects;
 pub mod railjson;
-
-pub use projects::{Ordering, Project};
+pub mod timetable;
+pub mod train_schedule;
 
 pub use documents::Document;
 pub use infra_objects::*;
+pub use projects::{Ordering, Project};
 
 use async_trait::async_trait;
 use diesel::{pg::Pg, result::Error::NotFound, AsChangeset, QueryableByName};

editoast/src/modelsv2/timetable.rs

@@ -0,0 +1,76 @@
+use crate::diesel::query_dsl::methods::DistinctDsl;
+use crate::error::Result;
+use crate::models::List;
+use crate::models::NoParams;
+use crate::modelsv2::{Retrieve, Row};
+use crate::tables::timetable_v2::dsl;
+use crate::views::pagination::Paginate;
+use crate::views::pagination::PaginatedResponse;
+use async_trait::async_trait;
+use diesel::sql_query;
+use diesel::sql_types::{Array, BigInt, Nullable};
+use diesel_async::AsyncPgConnection as PgConnection;
+use editoast_derive::ModelV2;
+
+#[derive(Debug, Default, Clone, ModelV2)]
+#[model(table = crate::tables::timetable_v2)]
+#[model(changeset(public))]
+pub struct Timetable {
+    pub id: i64,
+    pub electrical_profile_set_id: Option<i64>,
+}
+
+#[async_trait]
+impl List<NoParams> for Timetable {
+    async fn list_conn(
+        conn: &mut PgConnection,
+        page: i64,
+        page_size: i64,
+        _: NoParams,
+    ) -> Result<PaginatedResponse<Self>> {
+        let timetable_rows = dsl::timetable_v2
+            .distinct()
+            .paginate(page, page_size)
+            .load_and_count::<Row<Timetable>>(conn)
+            .await?;
+
+        Ok(timetable_rows.into())
+    }
+}
+
+/// Should be used to retrieve a timetable with its trains
+#[derive(Debug, Clone, QueryableByName)]
+pub struct TimetableWithTrains {
+    #[diesel(sql_type = BigInt)]
+    pub id: i64,
+    #[diesel(sql_type = Nullable<BigInt>)]
+    pub electrical_profile_set_id: Option<i64>,
+    #[diesel(sql_type = Array<BigInt>)]
+    pub train_ids: Vec<i64>,
+}
+
+#[async_trait::async_trait]
+impl Retrieve<i64> for TimetableWithTrains {
+    async fn retrieve(
+        conn: &mut diesel_async::AsyncPgConnection,
+        timetable_id: i64,
+    ) -> Result<Option<Self>> {
+        use diesel_async::RunQueryDsl;
+        let result = sql_query(
+            "SELECT timetable_v2.*,
+        array_remove(array_agg(train_schedule_v2.id), NULL) as train_ids
+        FROM timetable_v2
+        LEFT JOIN train_schedule_v2 ON timetable_v2.id = train_schedule_v2.timetable_id
+        WHERE timetable_v2.id = $1
+        GROUP BY timetable_v2.id",
+        )
+        .bind::<BigInt, _>(timetable_id)
+        .get_result::<TimetableWithTrains>(conn)
+        .await;
+        match result {
+            Ok(result) => Ok(Some(result)),
+            Err(diesel::result::Error::NotFound) => Ok(None),
+            Err(err) => Err(err.into()),
+        }
+    }
+}

editoast/src/modelsv2/train_schedule.rs

@@ -0,0 +1,30 @@
+use crate::schema::v2::trainschedule::{
+    Comfort, Distribution, Margins, PathItem, PowerRestrictionItem, ScheduleItem,
+    TrainScheduleOptions,
+};
+use crate::DieselJson;
+use chrono::{DateTime, Utc};
+use editoast_derive::ModelV2;
+
+#[derive(Debug, Default, Clone, ModelV2)]
+#[model(table = crate::tables::train_schedule_v2)]
+#[model(changeset(public))]
+pub struct TrainSchedule {
+    pub id: i64,
+    pub train_name: String,
+    pub labels: Vec<Option<String>>,
+    pub rolling_stock_name: String,
+    pub timetable_id: i64,
+    pub start_time: DateTime<Utc>,
+    pub schedule: DieselJson<Vec<ScheduleItem>>,
+    pub margins: DieselJson<Margins>,
+    pub initial_speed: f64,
+    #[model(to_enum)]
+    pub comfort: Comfort,
+    pub path: DieselJson<Vec<PathItem>>,
+    #[model(to_enum)]
+    pub constraint_distribution: Distribution,
+    pub speed_limit_tag: Option<String>,
+    pub power_restrictions: DieselJson<Vec<PowerRestrictionItem>>,
+    pub options: DieselJson<TrainScheduleOptions>,
+}