diff --git a/editoast/src/client/mod.rs b/editoast/src/client/mod.rs index 0fe0f3621c7..f4c45199886 100644 --- a/editoast/src/client/mod.rs +++ b/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, + pub path: PathBuf, } #[derive(Subcommand, Debug)] diff --git a/editoast/src/fixtures.rs b/editoast/src/fixtures.rs index 6092fd56736..6bf8d3ace89 100644 --- a/editoast/src/fixtures.rs +++ b/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, diff --git a/editoast/src/main.rs b/editoast/src/main.rs index 4d760945b96..20c60998026 100644 --- a/editoast/src/main.rs +++ b/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> { } 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, +) -> Result<(), Box> { + 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 = + serde_json::from_reader(BufReader::new(train_file))?; + let changesets: Vec = 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 { match (args.sentry_dsn.clone(), args.sentry_env.clone()) { (Some(sentry_dsn), Some(sentry_env)) => Some(sentry::init(( @@ -790,11 +853,13 @@ 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_async::RunQueryDsl; + use modelsv2::DeleteStatic; use rand::distributions::Alphanumeric; use rand::{thread_rng, Rng}; use rstest::rstest; @@ -802,6 +867,29 @@ mod tests { use std::io::Write; use tempfile::NamedTempFile; + #[rstest] + async fn import_train_schedule_v2(db_pool: Data) { + 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_all(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); + + Timetable::delete_static(conn, timetable.id).await.unwrap(); + } + #[rstest] async fn import_rolling_stock_ko_file_not_found(db_pool: Data) { // GIVEN diff --git a/editoast/src/tests/train_schedules/simple_array.json b/editoast/src/tests/train_schedules/simple_array.json new file mode 100644 index 00000000000..63a4479be69 --- /dev/null +++ b/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 + } + } +]