diff --git a/editoast/src/models/train_schedule.rs b/editoast/src/models/train_schedule.rs index 4b1890ce20c..4d4f849a01f 100644 --- a/editoast/src/models/train_schedule.rs +++ b/editoast/src/models/train_schedule.rs @@ -1,6 +1,8 @@ use chrono::DateTime; use chrono::Utc; use editoast_derive::Model; +use editoast_models::DbConnection; +use crate::error::Result; use editoast_schemas::train_schedule::Comfort; use editoast_schemas::train_schedule::Distribution; use editoast_schemas::train_schedule::Margins; @@ -9,8 +11,14 @@ use editoast_schemas::train_schedule::PowerRestrictionItem; use editoast_schemas::train_schedule::ScheduleItem; use editoast_schemas::train_schedule::TrainScheduleBase; use editoast_schemas::train_schedule::TrainScheduleOptions; +use diesel::prelude::*; +use diesel_async::RunQueryDsl; +use futures_util::stream::TryStreamExt; +use std::ops::DerefMut; use super::Model as _; +use crate::diesel::ExpressionMethods; +use crate::Row; #[derive(Debug, Default, Clone, Model)] #[model(table = editoast_models::tables::train_schedule)] @@ -74,3 +82,24 @@ impl From for TrainScheduleChangeset { .options(options) } } + +pub async fn filter_train_by_start_time_and_timetable( + conn: &mut DbConnection, + time: DateTime, + timetable_id: i64, +) -> Result> { + use editoast_models::tables::train_schedule::dsl; + let train_schedules = dsl::train_schedule + .filter(dsl::start_time.le(time)) + .filter(dsl::timetable_id.eq(timetable_id)) + .load_stream::>(conn.write().await.deref_mut()) + .await? + .map_ok(|ts| ts.into()) + .try_collect::>() + .await; + match train_schedules { + Ok(train_schedules) => Ok(train_schedules), + Err(err) => Err(err.into()), + + } +} diff --git a/editoast/src/views/timetable/stdcm.rs b/editoast/src/views/timetable/stdcm.rs index adcfef9de90..e6d1ff3f7f6 100644 --- a/editoast/src/views/timetable/stdcm.rs +++ b/editoast/src/views/timetable/stdcm.rs @@ -8,6 +8,7 @@ use axum::extract::State; use axum::Extension; use chrono::Utc; use chrono::{DateTime, Duration}; +use crate::models::train_schedule::filter_train_by_start_time_and_timetable; use editoast_authz::BuiltinRole; use editoast_derive::EditoastError; use editoast_models::DbConnectionPoolV2; @@ -40,7 +41,7 @@ use crate::core::CoreClient; use crate::error::Result; use crate::models::prelude::*; use crate::models::stdcm_log::StdcmLog; -use crate::models::timetable::TimetableWithTrains; +use crate::models::timetable::Timetable; use crate::models::train_schedule::TrainSchedule; use crate::models::Infra; use crate::models::RollingStockModel; @@ -133,6 +134,7 @@ async fn stdcm( Query(query): Query, Json(stdcm_request): Json, ) -> Result> { + let authorized = auth .check_roles([BuiltinRole::Stdcm].into()) .await @@ -148,19 +150,12 @@ async fn stdcm( let infra_id = query.infra; // 1. Retrieve Timetable / Infra / Trains / Simulation / Rolling Stock - let timetable_trains = TimetableWithTrains::retrieve_or_fail(&mut conn, timetable_id, || { - StdcmError::TimetableNotFound { timetable_id } - }) - .await?; let infra = Infra::retrieve_or_fail(&mut conn, infra_id, || StdcmError::InfraNotFound { infra_id, }) .await?; - let (train_schedules, _): (Vec<_>, _) = - TrainSchedule::retrieve_batch(&mut conn, timetable_trains.train_ids.clone()).await?; - let rolling_stock = RollingStockModel::retrieve_or_fail(&mut conn, stdcm_request.rolling_stock_id, || { StdcmError::RollingStockNotFound { @@ -191,6 +186,29 @@ async fn stdcm( let earliest_departure_time = stdcm_request.get_earliest_departure_time(simulation_run_time); let latest_simulation_end = stdcm_request.get_latest_simulation_end(simulation_run_time); + let timetable = Timetable::retrieve_or_fail(&mut conn, timetable_id, || { + StdcmError::TimetableNotFound { timetable_id } + }) + .await?; + // Filter train + let mut train_schedules = filter_train_by_start_time_and_timetable(&mut conn, latest_simulation_end, timetable.id).await?; + + train_schedules.retain(|train_schedule| { + if train_schedule.start_time >= earliest_departure_time { + true + } else { + if let Some(last_item) = train_schedule.schedule.last() { + if let Some(last_path_item) = train_schedule.path.last() { + last_item.at == last_path_item.id + } else { + false + } + } else { + false + } + } + }); + // 3. Get scheduled train requirements let simulations: Vec<_> = train_simulation_batch( &mut conn,