editoast: filter train schedules for LMR

Signed-off-by: Youness CHRIFI ALAOUI <[email protected]>

editoast/src/views/timetable/stdcm.rs

@@ -8,6 +8,8 @@ use axum::extract::State;
 use axum::Extension;
 use chrono::Utc;
 use chrono::{DateTime, Duration};
+use diesel::prelude::*;
+use diesel_async::RunQueryDsl;
 use editoast_authz::BuiltinRole;
 use editoast_derive::EditoastError;
 use editoast_models::DbConnectionPoolV2;
@@ -17,11 +19,13 @@ use editoast_schemas::train_schedule::Margins;
 use editoast_schemas::train_schedule::ReceptionSignal;
 use editoast_schemas::train_schedule::ScheduleItem;
 use failure_handler::SimulationFailureHandler;
+use futures_util::stream::TryStreamExt;
 use request::convert_steps;
 use request::Request;
 use serde::Deserialize;
 use serde::Serialize;
 use std::collections::HashMap;
+use std::ops::DerefMut;
 use std::sync::Arc;
 use thiserror::Error;
 use tracing::Instrument;
@@ -40,7 +44,6 @@ 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::train_schedule::TrainSchedule;
 use crate::models::Infra;
 use crate::models::RollingStockModel;
@@ -133,6 +136,8 @@ async fn stdcm(
     Query(query): Query<InfraIdQueryParam>,
     Json(stdcm_request): Json<Request>,
 ) -> Result<Json<StdcmResponse>> {
+    use editoast_models::tables::train_schedule::dsl;
+
     let authorized = auth
         .check_roles([BuiltinRole::Stdcm].into())
         .await
@@ -148,19 +153,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 +189,35 @@ 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 = dsl::train_schedule
+        .filter(dsl::start_time.ge(latest_simulation_end)).filter(dsl::timetable_id.eq(timetable.id))
+        .load_stream::<Row<TrainSchedule>>(conn.write().await.deref_mut())
+        .await?.map_ok(|ts| ts.into()).try_collect::<Vec<TrainSchedule>>().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,