Skip to content

Commit

Permalink
editoast: add train import command for train schedule v2
Browse files Browse the repository at this point in the history
  • Loading branch information
ElysaSrc committed Feb 28, 2024
1 parent a65a503 commit d815eeb
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 3 deletions.
16 changes: 16 additions & 0 deletions editoast/src/client/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down
56 changes: 53 additions & 3 deletions editoast/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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> {
Expand Down

0 comments on commit d815eeb

Please sign in to comment.