Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

editoast: add enabled window for stdcm search env table #10866

Merged
merged 1 commit into from
Feb 25, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions editoast/editoast_models/src/tables.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
ALTER TABLE stdcm_search_environment
DROP COLUMN enabled_from,
DROP COLUMN enabled_until;
Original file line number Diff line number Diff line change
@@ -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);
24 changes: 24 additions & 0 deletions editoast/openapi.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
15 changes: 10 additions & 5 deletions editoast/src/client/stdcm_search_env_commands.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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:#?}");

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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();
Expand Down
162 changes: 65 additions & 97 deletions editoast/src/models/stdcm_search_environment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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)
Expand All @@ -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;
Expand All @@ -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(
Expand All @@ -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,
Expand All @@ -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);
}
}
Loading
Loading