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 file")]
+pub struct ImportTrainArgs {
+    #[arg(long)]
+    pub timetable: Option<String>,
+    pub path: String,
 }
 
 #[derive(Subcommand, Debug)]

editoast/src/main.rs

@@ -21,7 +21,11 @@ use crate::core::CoreClient;
 use crate::error::InternalError;
 use crate::map::redis_utils::RedisClient;
 use crate::models::{Create, Delete, Infra};
+use crate::modelsv2::timetable::Timetable;
+use crate::modelsv2::Retrieve as RetrieveV2;
+use crate::modelsv2::{Create as CreateV2, Model};
 use crate::schema::electrical_profiles::ElectricalProfileSetData;
+use crate::schema::v2::trainschedule::TrainScheduleBase;
 use crate::schema::RailJson;
 use crate::views::infra::InfraForm;
 use crate::views::OpenApiRoot;
@@ -34,10 +38,13 @@ 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::train_schedule::TrainScheduleChangeset;
+use views::v2::train_schedule::TrainScheduleForm;
+
 use colored::*;
 use diesel::{sql_query, ConnectionError, ConnectionResult};
 use diesel_async::pooled_connection::deadpool::Pool;
@@ -184,7 +191,50 @@ 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 conn = &mut db_pool.get().await?;
+    let timetable = match args.timetable {
+        Some(timetable) => match Timetable::retrieve(conn, timetable.parse::<i64>()?).await? {
+            Some(timetable) => timetable,
+            None => {
+                let error = CliError::new(1, format!("❌ Timetable not found, Name: {timetable}"));
+                return Err(Box::new(error));
+            }
+        },
+        None => {
+            let changeset = Timetable::changeset();
+            changeset.create(conn).await?
+        }
+    };
+
+    let train_file = File::open(args.path)?;
+    let train_schedules: Vec<TrainScheduleBase> =
+        serde_json::from_reader(BufReader::new(train_file))?;
+    let ts_length = train_schedules.len();
+
+    for train_schedule in train_schedules {
+        let form = TrainScheduleForm {
+            timetable_id: timetable.id,
+            train_schedule,
+        };
+        let changeset: TrainScheduleChangeset = form.into();
+        changeset.create(conn).await?;
     }
+
+    println!(
+        "✅ {} train schedules created for timetable with id {}", ts_length, timetable.id
+    );
+
+    Ok(())
 }
 
 fn init_sentry(args: &RunserverArgs) -> Option<ClientInitGuard> {