editoast: add enabled window for stdcm search env table

Signed-off-by: Florian Amsallem <[email protected]>
OpenRailAssociation · Feb 25, 2025 · 57034a3 · 57034a3
1 parent b53ba93
commit 57034a3
Show file tree

Hide file tree

Showing 9 changed files with 181 additions and 134 deletions.
diff --git a/editoast/editoast_models/src/tables.rs b/editoast/editoast_models/src/tables.rs
@@ -671,6 +671,8 @@ diesel::table! {
         search_window_begin -> Timestamptz,
         search_window_end -> Timestamptz,
         temporary_speed_limit_group_id -> Nullable<Int8>,
+        enabled_from -> Timestamptz,
+        enabled_until -> Timestamptz,
     }
 }
 

diff --git a/...oast/migrations/2025-02-19-173042_adapt_search_environment_adding_enabled_window/down.sql b/...oast/migrations/2025-02-19-173042_adapt_search_environment_adding_enabled_window/down.sql
@@ -0,0 +1,3 @@
+ALTER TABLE stdcm_search_environment
+    DROP COLUMN enabled_from,
+    DROP COLUMN enabled_until;
diff --git a/editoast/migrations/2025-02-19-173042_adapt_search_environment_adding_enabled_window/up.sql b/editoast/migrations/2025-02-19-173042_adapt_search_environment_adding_enabled_window/up.sql
@@ -0,0 +1,5 @@
+ALTER TABLE stdcm_search_environment
+    ADD COLUMN enabled_from TIMESTAMP WITH TIME ZONE DEFAULT '1970-01-01 00:00:00' NOT NULL,
+    ADD COLUMN enabled_until TIMESTAMP WITH TIME ZONE DEFAULT '1970-01-02 00:00:00' NOT NULL;
+CREATE INDEX idx_search_environment_enabled_from ON stdcm_search_environment (enabled_from);
+CREATE INDEX idx_search_environment_enabled_until ON stdcm_search_environment (enabled_until);
diff --git a/editoast/openapi.yaml b/editoast/openapi.yaml
@@ -11372,10 +11372,22 @@ components:
       - timetable_id
       - search_window_begin
       - search_window_end
+      - enabled_from
+      - enabled_until
       properties:
         electrical_profile_set_id:
           type: integer
           format: int64
+        enabled_from:
+          type: string
+          format: date-time
+          description: The time window start point where the environment is enabled.
+        enabled_until:
+          type: string
+          format: date-time
+          description: |-
+            The time window end point where the environment is enabled.
+            This value is usually lower than the `search_window_begin`, since a search is performed before the train rolls.
         id:
           type: integer
           format: int64
@@ -11385,9 +11397,13 @@ components:
         search_window_begin:
           type: string
           format: date-time
+          description: |-
+            The start of the search time window.
+            Usually, trains schedules from the `timetable_id` runs within this window.
         search_window_end:
           type: string
           format: date-time
+          description: The end of the search time window.
         temporary_speed_limit_group_id:
           type: integer
           format: int64
@@ -11404,11 +11420,19 @@ components:
       - timetable_id
       - search_window_begin
       - search_window_end
+      - enabled_from
+      - enabled_until
       properties:
         electrical_profile_set_id:
           type: integer
           format: int64
           nullable: true
+        enabled_from:
+          type: string
+          format: date-time
+        enabled_until:
+          type: string
+          format: date-time
         infra_id:
           type: integer
           format: int64

diff --git a/editoast/src/client/stdcm_search_env_commands.rs b/editoast/src/client/stdcm_search_env_commands.rs
@@ -2,6 +2,7 @@ use crate::models::electrical_profiles::ElectricalProfileSet;
 use crate::models::stdcm_search_environment::StdcmSearchEnvironment;
 use crate::models::timetable::Timetable;
 use crate::models::work_schedules::WorkScheduleGroup;
+use crate::models::Create;
 use crate::models::Infra;
 use crate::models::Scenario;
 use crate::CliError;
@@ -103,7 +104,9 @@ async fn set_stdcm_search_env_from_scenario(
         .timetable_id(scenario.timetable_id)
         .search_window_begin(begin)
         .search_window_end(end)
-        .overwrite(conn)
+        .enabled_from(Utc::now())
+        .enabled_until(Utc::now() + Duration::days(1000))
+        .create(conn)
         .await?;
 
     println!("✅ STDCM search environment set up successfully");
@@ -170,7 +173,9 @@ async fn set_stdcm_search_env_from_scratch(
         .timetable_id(args.timetable_id)
         .search_window_begin(begin)
         .search_window_end(end)
-        .overwrite(conn)
+        .enabled_from(Utc::now())
+        .enabled_until(Utc::now() + Duration::days(1000))
+        .create(conn)
         .await?;
 
     println!("✅ STDCM search environment set up successfully");
@@ -213,7 +218,7 @@ async fn resolve_search_window(
 async fn show_stdcm_search_env(
     conn: &mut DbConnection,
 ) -> Result<(), Box<dyn Error + Send + Sync>> {
-    let search_env = StdcmSearchEnvironment::retrieve_latest(conn).await;
+    let search_env = StdcmSearchEnvironment::retrieve_latest_enabled(conn).await;
     if let Some(search_env) = search_env {
         println!("{search_env:#?}");
 
@@ -389,7 +394,7 @@ mod tests {
         let result = set_stdcm_search_env_from_scenario(args, conn).await;
         assert!(result.is_ok());
 
-        let search_env = StdcmSearchEnvironment::retrieve_latest(conn).await;
+        let search_env = StdcmSearchEnvironment::retrieve_latest_enabled(conn).await;
 
         assert!(search_env.is_some());
         let search_env = search_env.unwrap();
@@ -433,7 +438,7 @@ mod tests {
         let result = set_stdcm_search_env_from_scratch(args, conn).await;
         assert!(result.is_ok());
 
-        let search_env = StdcmSearchEnvironment::retrieve_latest(conn).await;
+        let search_env = StdcmSearchEnvironment::retrieve_latest_enabled(conn).await;
 
         assert!(search_env.is_some());
         let search_env = search_env.unwrap();

diff --git a/editoast/src/models/stdcm_search_environment.rs b/editoast/src/models/stdcm_search_environment.rs
@@ -4,15 +4,15 @@ use diesel::ExpressionMethods;
 use diesel::QueryDsl;
 use diesel_async::RunQueryDsl;
 use editoast_derive::Model;
-use editoast_models::model;
 use editoast_models::DbConnection;
 use serde::Serialize;
 use std::ops::DerefMut;
-use std::result::Result;
 use utoipa::ToSchema;
 
 use crate::models::prelude::*;
 
+#[cfg(test)]
+use editoast_models::model;
 #[cfg(test)]
 use serde::Deserialize;
 
@@ -30,25 +30,51 @@ pub struct StdcmSearchEnvironment {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub work_schedule_group_id: Option<i64>,
     pub timetable_id: i64,
+    /// The start of the search time window.
+    /// Usually, trains schedules from the `timetable_id` runs within this window.
     pub search_window_begin: DateTime<Utc>,
+    /// The end of the search time window.
     pub search_window_end: DateTime<Utc>,
     #[schema(nullable = false)]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub temporary_speed_limit_group_id: Option<i64>,
+    /// The time window start point where the environment is enabled.
+    pub enabled_from: DateTime<Utc>,
+    /// The time window end point where the environment is enabled.
+    /// This value is usually lower than the `search_window_begin`, since a search is performed before the train rolls.
+    pub enabled_until: DateTime<Utc>,
 }
 
 impl StdcmSearchEnvironment {
-    /// Retrieve the latest search environment. Returns None if no search environment is found.
-    pub async fn retrieve_latest(conn: &mut DbConnection) -> Option<Self> {
+    /// Retrieve the enabled search environment. If no env is enabled returns the most recent `enabled_until`.
+    /// In case of multiple enabled environments, the one with the highest `id` is returned.
+    pub async fn retrieve_latest_enabled(conn: &mut DbConnection) -> Option<Self> {
         use editoast_models::tables::stdcm_search_environment::dsl::*;
+        // Search for enabled env
+        let enabled_env = stdcm_search_environment
+            .order_by(id.desc())
+            .filter(enabled_from.le(diesel::dsl::now))
+            .filter(enabled_until.ge(diesel::dsl::now))
+            .first::<Row<StdcmSearchEnvironment>>(conn.write().await.deref_mut())
+            .await
+            .map(Into::into)
+            .ok();
+        if enabled_env.is_some() {
+            return enabled_env;
+        }
+
+        // Search for the most recent env
+        tracing::warn!("No STDCM search environment enabled");
         stdcm_search_environment
-            .order_by((search_window_end.desc(), search_window_begin.asc()))
+            .order_by((enabled_until.desc(), id.desc()))
             .first::<Row<StdcmSearchEnvironment>>(conn.write().await.deref_mut())
             .await
             .map(Into::into)
             .ok()
     }
 
+    /// Delete all existing search environments.
+    #[cfg(test)]
     pub async fn delete_all(conn: &mut DbConnection) -> Result<(), model::Error> {
         use editoast_models::tables::stdcm_search_environment::dsl::*;
         diesel::delete(stdcm_search_environment)
@@ -58,18 +84,10 @@ impl StdcmSearchEnvironment {
     }
 }
 
-impl StdcmSearchEnvironmentChangeset {
-    pub async fn overwrite(
-        self,
-        conn: &mut DbConnection,
-    ) -> Result<StdcmSearchEnvironment, model::Error> {
-        StdcmSearchEnvironment::delete_all(conn).await?;
-        self.create(conn).await
-    }
-}
-
 #[cfg(test)]
 pub mod tests {
+    use chrono::Duration;
+    use chrono::DurationRound;
     use chrono::TimeZone;
     use chrono::Utc;
     use pretty_assertions::assert_eq;
@@ -85,7 +103,7 @@ pub mod tests {
     use crate::models::timetable::Timetable;
     use crate::models::work_schedules::WorkScheduleGroup;
     use crate::models::Infra;
-    use crate::models::{Count, Model, SelectionSettings};
+    use crate::models::Model;
     use editoast_models::DbConnectionPoolV2;
 
     pub async fn stdcm_search_env_fixtures(
@@ -112,76 +130,10 @@ pub mod tests {
         )
     }
 
-    #[rstest]
-    async fn test_overwrite() {
-        let db_pool = DbConnectionPoolV2::for_tests();
-        let initial_env_count =
-            StdcmSearchEnvironment::count(&mut db_pool.get_ok(), Default::default())
-                .await
-                .expect("failed to count STDCM envs");
-        let (
-            infra,
-            timetable,
-            work_schedule_group,
-            temporary_speed_limit_group,
-            electrical_profile_set,
-        ) = stdcm_search_env_fixtures(&mut db_pool.get_ok()).await;
-
-        let changeset_1 = StdcmSearchEnvironment::changeset()
-            .infra_id(infra.id)
-            .electrical_profile_set_id(Some(electrical_profile_set.id))
-            .work_schedule_group_id(Some(work_schedule_group.id))
-            .temporary_speed_limit_group_id(Some(temporary_speed_limit_group.id))
-            .timetable_id(timetable.id)
-            .search_window_begin(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap())
-            .search_window_end(Utc.with_ymd_and_hms(2024, 1, 15, 0, 0, 0).unwrap());
-
-        let begin = Utc.with_ymd_and_hms(2024, 1, 16, 0, 0, 0).unwrap();
-        let end = Utc.with_ymd_and_hms(2024, 1, 13, 0, 0, 0).unwrap();
-
-        let changeset_2 = changeset_1
-            .clone()
-            .search_window_begin(begin)
-            .search_window_end(end);
-
-        changeset_1
-            .create(&mut db_pool.get_ok())
-            .await
-            .expect("Failed to create first search environment");
-
-        assert_eq!(
-            StdcmSearchEnvironment::count(&mut db_pool.get_ok(), SelectionSettings::new())
-                .await
-                .expect("Failed to count"),
-            initial_env_count + 1
-        );
-
-        let _ = changeset_2
-            .overwrite(&mut db_pool.get_ok())
-            .await
-            .expect("Failed to overwrite search environment");
-
-        assert_eq!(
-            StdcmSearchEnvironment::count(&mut db_pool.get_ok(), SelectionSettings::new())
-                .await
-                .expect("Failed to count"),
-            1
-        );
-
-        let result = StdcmSearchEnvironment::retrieve_latest(&mut db_pool.get_ok())
-            .await
-            .expect("Failed to retrieve latest search environment");
-
-        assert_eq!(result.search_window_begin, begin);
-        assert_eq!(result.search_window_end, end);
-    }
-
     #[rstest]
     async fn test_retrieve_latest() {
         let db_pool = DbConnectionPoolV2::for_tests();
-        StdcmSearchEnvironment::delete_all(&mut db_pool.get_ok())
-            .await
-            .expect("failed to delete envs");
+
         let (
             infra,
             timetable,
@@ -197,43 +149,59 @@ pub mod tests {
             .temporary_speed_limit_group_id(Some(temporary_speed_limit_group.id))
             .timetable_id(timetable.id)
             .search_window_begin(Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap())
-            .search_window_end(Utc.with_ymd_and_hms(2024, 1, 15, 0, 0, 0).unwrap());
+            .search_window_end(Utc.with_ymd_and_hms(2024, 1, 2, 0, 0, 0).unwrap())
+            .enabled_from(Utc::now() - Duration::days(3))
+            .enabled_until(Utc::now() - Duration::days(2));
 
         let too_young = too_old
             .clone()
-            .search_window_begin(Utc.with_ymd_and_hms(2024, 1, 16, 0, 0, 0).unwrap())
-            .search_window_end(Utc.with_ymd_and_hms(2024, 1, 31, 0, 0, 0).unwrap());
+            .enabled_from(Utc::now() + Duration::days(2))
+            .enabled_until(Utc::now() + Duration::days(3));
+
+        let enabled_but_not_last = too_old
+            .clone()
+            .enabled_from(Utc::now() - Duration::hours(1))
+            .enabled_until(Utc::now() + Duration::hours(1));
 
-        let begin = Utc.with_ymd_and_hms(2024, 1, 7, 0, 0, 0).unwrap();
-        let end = Utc.with_ymd_and_hms(2024, 1, 31, 0, 0, 0).unwrap();
+        let enabled_from =
+            Utc::now().duration_trunc(Duration::seconds(1)).unwrap() - Duration::days(1);
+        let enabled_until =
+            Utc::now().duration_trunc(Duration::seconds(1)).unwrap() + Duration::days(1);
 
         let the_best = too_old
             .clone()
-            .search_window_begin(begin)
-            .search_window_end(end);
-
-        for changeset in [too_old, too_young.clone(), the_best, too_young] {
+            .enabled_from(enabled_from)
+            .enabled_until(enabled_until);
+
+        for changeset in [
+            too_old,
+            too_young.clone(),
+            enabled_but_not_last,
+            the_best,
+            too_young,
+        ] {
             changeset
                 .create(&mut db_pool.get_ok())
                 .await
                 .expect("Failed to create search environment");
         }
 
-        let result = StdcmSearchEnvironment::retrieve_latest(&mut db_pool.get_ok())
+        let result = StdcmSearchEnvironment::retrieve_latest_enabled(&mut db_pool.get_ok())
             .await
             .expect("Failed to retrieve latest search environment");
 
-        assert_eq!(result.search_window_begin, begin);
-        assert_eq!(result.search_window_end, end);
+        assert_eq!(result.enabled_from, enabled_from);
+        assert_eq!(result.enabled_until, enabled_until);
     }
 
     #[rstest]
     async fn test_retrieve_latest_empty() {
         let db_pool = DbConnectionPoolV2::for_tests();
         StdcmSearchEnvironment::delete_all(&mut db_pool.get_ok())
             .await
-            .expect("failed to delete envs");
-        let result = StdcmSearchEnvironment::retrieve_latest(&mut db_pool.get_ok()).await;
+            .expect("Failed to delete all search environments");
+
+        let result = StdcmSearchEnvironment::retrieve_latest_enabled(&mut db_pool.get_ok()).await;
         assert_eq!(result, None);
     }
 }