editoast: add train import command for train schedule v2

Author: ElysaSrc

editoast/src/client/mod.rs

@@ -50,6 +50,22 @@ pub enum Commands {
     Search(SearchCommands),
     #[command(subcommand, about, long_about = "Infrastructure related commands")]
     Infra(InfraCommands),
+    #[command(subcommand, about, long_about = "Trains related commands")]
+    Trains(TrainsCommands),
+}
+
+#[derive(Subcommand, Debug)]
+pub enum TrainsCommands {
+    Import(ImportTrainArgs),
+}
+
+#[derive(Args, Debug, Derivative)]
+#[derivative(Default)]
+#[command(about, long_about = "Import a train given a JSON file")]
+pub struct ImportTrainArgs {
+    #[arg(long, help = "The timetable id on which attach the trains to")]
+    pub timetable: Option<i64>,
+    pub path: PathBuf,
 }
 
 #[derive(Subcommand, Debug)]

editoast/src/fixtures.rs

@@ -142,6 +142,10 @@ pub mod tests {
         rs
     }
 
+    pub fn get_trainschedule_json_array() -> &'static str {
+        include_str!("./tests/train_schedules/simple_array.json")
+    }
+
     pub async fn named_other_rolling_stock(
         name: &str,
         db_pool: Data<DbPool>,

editoast/src/main.rs

@@ -27,10 +27,17 @@ use chashmap::CHashMap;
 use clap::Parser;
 use client::{
     ClearArgs, Client, Color, Commands, DeleteProfileSetArgs, ElectricalProfilesCommands,
-    GenerateArgs, ImportProfileSetArgs, ImportRailjsonArgs, ImportRollingStockArgs, InfraCloneArgs,
-    InfraCommands, ListProfileSetArgs, MakeMigrationArgs, RedisConfig, RefreshArgs, RunserverArgs,
-    SearchCommands,
+    GenerateArgs, ImportProfileSetArgs, ImportRailjsonArgs, ImportRollingStockArgs,
+    ImportTrainArgs, InfraCloneArgs, InfraCommands, ListProfileSetArgs, MakeMigrationArgs,
+    RedisConfig, RefreshArgs, RunserverArgs, SearchCommands, TrainsCommands,
 };
+use modelsv2::{
+    timetable::Timetable, train_schedule::TrainSchedule, train_schedule::TrainScheduleChangeset,
+    Create as CreateV2, CreateBatch, Model, Retrieve as RetrieveV2,
+};
+use schema::v2::trainschedule::TrainScheduleBase;
+use views::v2::train_schedule::TrainScheduleForm;
+
 use colored::*;
 use core::CoreClient;
 use diesel::{sql_query, ConnectionError, ConnectionResult};
@@ -185,9 +192,65 @@ async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
             }
             InfraCommands::ImportRailjson(args) => import_railjson(args, create_db_pool()?).await,
         },
+        Commands::Trains(subcommand) => match subcommand {
+            TrainsCommands::Import(args) => trains_import(args, create_db_pool()?).await,
+        },
     }
 }
 
+async fn trains_import(
+    args: ImportTrainArgs,
+    db_pool: Data<DbPool>,
+) -> Result<(), Box<dyn Error + Send + Sync>> {
+    let train_file = match File::open(args.path.clone()) {
+        Ok(file) => file,
+        Err(e) => {
+            let error = CliError::new(
+                1,
+                format!("❌ Could not open file {:?} ({:?})", args.path, e),
+            );
+            return Err(Box::new(error));
+        }
+    };
+
+    let conn = &mut db_pool.get().await?;
+    let timetable = match args.timetable {
+        Some(timetable) => match Timetable::retrieve(conn, timetable).await? {
+            Some(timetable) => timetable,
+            None => {
+                let error = CliError::new(1, format!("❌ Timetable not found, id: {0}", timetable));
+                return Err(Box::new(error));
+            }
+        },
+        None => {
+            let changeset = Timetable::changeset();
+            changeset.create(conn).await?
+        }
+    };
+
+    let train_schedules: Vec<TrainScheduleBase> =
+        serde_json::from_reader(BufReader::new(train_file))?;
+    let changesets: Vec<TrainScheduleChangeset> = train_schedules
+        .into_iter()
+        .map(|train_schedule| {
+            TrainScheduleForm {
+                timetable_id: timetable.id,
+                train_schedule,
+            }
+            .into()
+        })
+        .collect();
+    let inserted: Vec<_> = TrainSchedule::create_batch(conn, changesets).await?;
+
+    println!(
+        "✅ {} train schedules created for timetable with id {}",
+        inserted.len(),
+        timetable.id
+    );
+
+    Ok(())
+}
+
 fn init_sentry(args: &RunserverArgs) -> Option<ClientInitGuard> {
     match (args.sentry_dsn.clone(), args.sentry_env.clone()) {
         (Some(sentry_dsn), Some(sentry_env)) => Some(sentry::init((
@@ -790,10 +853,11 @@ mod tests {
     use super::*;
 
     use crate::fixtures::tests::{
-        db_pool, electrical_profile_set, get_fast_rolling_stock, TestFixture,
+        db_pool, electrical_profile_set, get_fast_rolling_stock, get_trainschedule_json_array,
+        TestFixture,
     };
     use diesel::sql_query;
-    use diesel::sql_types::Text;
+    use diesel::sql_types::{BigInt, Text};
     use diesel_async::RunQueryDsl;
     use rand::distributions::Alphanumeric;
     use rand::{thread_rng, Rng};
@@ -802,6 +866,33 @@ mod tests {
     use std::io::Write;
     use tempfile::NamedTempFile;
 
+    #[rstest]
+    async fn import_train_schedule_v2(db_pool: Data<DbPool>) {
+        let conn = &mut db_pool.get().await.unwrap();
+
+        let changeset = Timetable::changeset();
+        let timetable = changeset.create(conn).await.unwrap();
+
+        let mut file = NamedTempFile::new().unwrap();
+        file.write(get_trainschedule_json_array().as_bytes())
+            .unwrap();
+
+        let args = ImportTrainArgs {
+            path: file.path().into(),
+            timetable: Some(timetable.id),
+        };
+
+        let result = trains_import(args, db_pool.clone()).await;
+
+        assert!(result.is_ok(), "{:?}", result);
+
+        sql_query("DELETE FROM timetable_v2 WHERE id = $1")
+            .bind::<BigInt, _>(timetable.id)
+            .execute(conn)
+            .await
+            .unwrap();
+    }
+
     #[rstest]
     async fn import_rolling_stock_ko_file_not_found(db_pool: Data<DbPool>) {
         // GIVEN

editoast/src/tests/train_schedules/simple_array.json

@@ -0,0 +1,77 @@
+[
+    {
+        "train_name": "ABC3615",
+        "rolling_stock_name": "R2D2",
+        "labels": [
+          "choo-choo",
+          "tchou-tchou"
+        ],
+        "speed_limit_tag": "MA100",
+        "start_time": "2023-12-21T08:51:30+00:00",
+        "path": [
+          {
+            "id": "a",
+            "uic": 87210
+          },
+          {
+            "id": "b",
+            "track": "foo",
+            "offset": 10
+          },
+          {
+            "id": "c",
+            "deleted": true,
+            "trigram": "ABC"
+          },
+          {
+            "id": "d",
+            "operational_point": "X"
+          }
+        ],
+        "constraint_distribution": "MARECO",
+        "schedule": [
+          {
+            "at": "a",
+            "stop_for": "PT5M",
+            "locked": true
+          },
+          {
+            "at": "b",
+            "arrival": "PT10M",
+            "stop_for": "PT5M"
+          },
+          {
+            "at": "c",
+            "stop_for": "PT5M"
+          },
+          {
+            "at": "d",
+            "arrival": "PT50M",
+            "locked": true
+          }
+        ],
+        "margins": {
+          "boundaries": [
+            "b",
+            "c"
+          ],
+          "values": [
+            "5%",
+            "3min/km",
+            "none"
+          ]
+        },
+        "initial_speed": 2.5,
+        "power_restrictions": [
+          {
+            "from": "b",
+            "to": "c",
+            "value": "M1C1"
+          }
+        ],
+        "comfort": "AIR_CONDITIONING",
+        "options": {
+          "use_electrical_profiles": true
+        }
+    }      
+]