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, id: {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> {

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
+        }
+    }      
+]