Adding OperationalPointIdLocation as location for POST /timetable/{id}

Author: bloussou

editoast/src/models/infra_objects/operational_point.rs

@@ -3,7 +3,7 @@
 use crate::error::Result;
 use crate::{schema::OperationalPoint, tables::infra_object_operational_point};
 use derivative::Derivative;
-use diesel::sql_types::{Array, BigInt};
+use diesel::sql_types::{Array, BigInt, Text};
 use diesel::{result::Error as DieselError, sql_query};
 use diesel::{ExpressionMethods, QueryDsl};
 use diesel_async::{AsyncPgConnection as PgConnection, RunQueryDsl};
@@ -38,4 +38,19 @@ impl OperationalPointModel {
             .load(conn)
             .await?)
     }
+
+    /// Retrieve a list of operational points from the database
+    pub async fn retrieve_from_ids(
+        conn: &mut PgConnection,
+        infra_id: i64,
+        ids: &[String],
+    ) -> Result<Vec<Self>> {
+        let query = "SELECT * FROM infra_object_operational_point
+                                WHERE infra_id = $1 AND infra_object_operational_point.obj_id = ANY($2)".to_string();
+        Ok(sql_query(query)
+            .bind::<BigInt, _>(infra_id)
+            .bind::<Array<Text>, _>(ids)
+            .load(conn)
+            .await?)
+    }
 }

editoast/src/views/timetable/import.rs

@@ -67,6 +67,8 @@ pub enum TimetableImportPathLocation {
     TrackOffsetLocation { track_section: String, offset: f64 },
     #[serde(rename = "operational_point")]
     OperationalPointLocation { uic: i64 },
+    #[serde(rename = "operational_point_id")]
+    OperationalPointIdLocation { id: String },
 }
 
 #[derive(Debug, Serialize, ToSchema)]
@@ -77,7 +79,7 @@ struct TimetableImportReport {
 #[derive(Debug, Clone, Serialize, ToSchema)]
 enum TimetableImportError {
     RollingStockNotFound { name: String },
-    OperationalPointNotFound { missing_uics: Vec<i64> },
+    OperationalPointNotFound { missing_uics: Vec<i64>, missing_ids: Vec<String> },
     PathfindingError { cause: InternalError },
     SimulationError { cause: InternalError },
 }
@@ -183,9 +185,10 @@ async fn import_item(
     pf_request.with_rolling_stocks(&mut vec![rolling_stock.clone()]);
     // List operational points uic needed for this import
     let ops_uic = ops_uic_from_path(&import_item.path);
+    let ops_id = ops_id_from_path(&import_item.path);
     // Retrieve operational points
-    let op_id_to_parts = match find_operation_points(&ops_uic, infra_id, conn).await? {
-        Ok(op_id_to_parts) => op_id_to_parts,
+    let (op_uic_to_parts, op_id_to_parts) = match find_operation_points(&ops_uic, &ops_id, infra_id, conn).await? {
+        Ok((op_uic_to_parts, op_id_to_parts)) => (op_uic_to_parts, op_id_to_parts),
         Err(err) => {
             return Ok(import_item
                 .trains
@@ -195,7 +198,7 @@ async fn import_item(
         }
     };
     // Create waypoints
-    let mut waypoints = waypoints_from_steps(&import_item.path, &op_id_to_parts);
+    let mut waypoints = waypoints_from_steps(&import_item.path, &op_uic_to_parts, &op_id_to_parts);
     pf_request.with_waypoints(&mut waypoints);
 
     // Run pathfinding
@@ -295,30 +298,50 @@ async fn import_item(
 
 async fn find_operation_points(
     ops_uic: &[i64],
+    ops_id: &[String],
     infra_id: i64,
     conn: &mut AsyncPgConnection,
-) -> Result<std::result::Result<HashMap<i64, Vec<OperationalPointPart>>, TimetableImportError>> {
+) -> Result<std::result::Result<(HashMap<i64, Vec<OperationalPointPart>>, HashMap<String, Vec<OperationalPointPart>>), TimetableImportError>> {
     // Retrieve operational points
     let ops = OperationalPointModel::retrieve_from_uic(conn, infra_id, ops_uic).await?;
-    let mut op_id_to_parts = HashMap::<_, Vec<_>>::new();
+    let mut op_uic_to_parts = HashMap::<_, Vec<_>>::new();
     for op in ops {
-        op_id_to_parts
+        op_uic_to_parts
             .entry(op.data.0.extensions.identifier.unwrap().uic)
             .or_default()
             .extend(op.data.0.parts);
     }
+    let mut missing_uics: Vec<i64> = vec![];
     // If we didn't find all the operational points, we can't run the pathfinding
-    if op_id_to_parts.len() != ops_uic.len() {
-        let missing_uics = ops_uic
+    if op_uic_to_parts.len() != ops_uic.len() {
+        ops_uic
             .iter()
-            .filter(|uic| !op_id_to_parts.contains_key(uic))
-            .cloned()
-            .collect();
+            .for_each(|uic| {if !op_uic_to_parts.contains_key(uic) {missing_uics.push(*uic)}});
+    }
+
+    let ops2 = OperationalPointModel::retrieve_from_ids(conn, infra_id, ops_id).await?;
+    let mut op_id_to_parts = HashMap::<_, Vec<_>>::new();
+    for op in ops2 {
+        op_id_to_parts
+            .entry(op.obj_id)
+            .or_default()
+            .extend(op.data.0.parts);
+    }
+    // If we didn't find all the operational points, we can't run the pathfinding
+    let mut missing_ids: Vec<String> = vec![];
+    if op_id_to_parts.len() != ops_id.len() {
+        ops_id
+            .iter()
+            .for_each(|id| {if !op_id_to_parts.contains_key(id) {missing_ids.push(id.to_string())}});
+    }
+    if missing_uics.len() + missing_ids.len() > 0 {
         return Ok(Err(TimetableImportError::OperationalPointNotFound {
             missing_uics,
+            missing_ids,
         }));
     }
-    Ok(Ok(op_id_to_parts))
+
+    Ok(Ok((op_uic_to_parts, op_id_to_parts)))
 }
 
 fn ops_uic_from_path(path: &[TimetableImportPathStep]) -> Vec<i64> {
@@ -335,9 +358,24 @@ fn ops_uic_from_path(path: &[TimetableImportPathStep]) -> Vec<i64> {
     ops_uic
 }
 
+fn ops_id_from_path(path: &[TimetableImportPathStep]) -> Vec<String> {
+    let mut ops_id = path
+        .iter()
+        .filter_map(|step| match &step.location {
+            TimetableImportPathLocation::OperationalPointIdLocation { id } => Some(id.to_string()),
+            _ => None,
+        })
+        .collect::<Vec<_>>();
+    // Remove duplicates
+    ops_id.sort();
+    ops_id.dedup();
+    ops_id
+}
+
 fn waypoints_from_steps(
     path: &Vec<TimetableImportPathStep>,
-    op_id_to_parts: &HashMap<i64, Vec<OperationalPointPart>>,
+    op_uic_to_parts: &HashMap<i64, Vec<OperationalPointPart>>,
+    op_id_to_parts: &HashMap<String, Vec<OperationalPointPart>>,
 ) -> Vec<Vec<Waypoint>> {
     let mut res = PathfindingWaypoints::new();
     for step in path {
@@ -346,12 +384,18 @@ fn waypoints_from_steps(
                 track_section,
                 offset,
             } => Vec::from(Waypoint::bidirectional(track_section, *offset)),
-            TimetableImportPathLocation::OperationalPointLocation { uic } => op_id_to_parts
+            TimetableImportPathLocation::OperationalPointLocation { uic } => op_uic_to_parts
                 .get(uic)
                 .unwrap()
                 .iter()
                 .flat_map(|op_part| Waypoint::bidirectional(&op_part.track, op_part.position))
                 .collect(),
+            TimetableImportPathLocation::OperationalPointIdLocation { id } => op_id_to_parts
+                .get(id)
+                .unwrap()
+                .iter()
+                .flat_map(|op_part| Waypoint::bidirectional(&op_part.track, op_part.position))
+                .collect(),
         });
     }
     res
@@ -438,8 +482,8 @@ mod tests {
 
     #[test]
     fn test_waypoints_from_steps() {
-        let mut op_id_to_parts = HashMap::new();
-        op_id_to_parts.insert(
+        let mut op_uic_to_parts = HashMap::new();
+        op_uic_to_parts.insert(
             1,
             vec![
                 OperationalPointPart {
@@ -455,6 +499,23 @@ mod tests {
             ],
         );
 
+        let mut op_id_to_parts = HashMap::new();
+        op_id_to_parts.insert(
+            "a1".to_string(),
+            vec![
+                OperationalPointPart {
+                    track: Identifier("E".to_string()),
+                    position: 0.,
+                    ..Default::default()
+                },
+                OperationalPointPart {
+                    track: Identifier("F".to_string()),
+                    position: 100.,
+                    ..Default::default()
+                },
+            ],
+        );
+
         let path = vec![
             TimetableImportPathStep {
                 location: TimetableImportPathLocation::TrackOffsetLocation {
@@ -467,11 +528,15 @@ mod tests {
                 location: TimetableImportPathLocation::OperationalPointLocation { uic: 1 },
                 schedule: HashMap::new(),
             },
+            TimetableImportPathStep {
+                location: TimetableImportPathLocation::OperationalPointIdLocation { id: "a1".to_string() },
+                schedule: HashMap::new(),
+            },
         ];
 
-        let waypoints = waypoints_from_steps(&path, &op_id_to_parts);
+        let waypoints = waypoints_from_steps(&path, &op_uic_to_parts, &op_id_to_parts);
 
-        assert_eq!(waypoints.len(), 2);
+        assert_eq!(waypoints.len(), 3);
         assert_eq!(waypoints[0], Waypoint::bidirectional("C", 50.));
         assert_eq!(
             waypoints[1],
@@ -481,10 +546,18 @@ mod tests {
             ]
             .concat()
         );
+        assert_eq!(
+            waypoints[2],
+            [
+                Waypoint::bidirectional("E", 0.),
+                Waypoint::bidirectional("F", 100.),
+            ]
+            .concat()
+        );
     }
 
     #[test]
-    fn test_ops_uic_from_path() {
+    fn test_ops_uic_id_from_path() {
         let path = vec![
             TimetableImportPathStep {
                 location: TimetableImportPathLocation::TrackOffsetLocation {
@@ -497,6 +570,10 @@ mod tests {
                 location: TimetableImportPathLocation::OperationalPointLocation { uic: 1 },
                 schedule: HashMap::new(),
             },
+            TimetableImportPathStep {
+                location: TimetableImportPathLocation::OperationalPointIdLocation { id: "a1".to_string() },
+                schedule: HashMap::new(),
+            },
             TimetableImportPathStep {
                 location: TimetableImportPathLocation::OperationalPointLocation { uic: 2 },
                 schedule: HashMap::new(),
@@ -508,6 +585,14 @@ mod tests {
                 },
                 schedule: HashMap::new(),
             },
+            TimetableImportPathStep {
+                location: TimetableImportPathLocation::OperationalPointIdLocation { id: "a2".to_string() },
+                schedule: HashMap::new(),
+            },
+            TimetableImportPathStep {
+                location: TimetableImportPathLocation::OperationalPointIdLocation { id: "a1".to_string() },
+                schedule: HashMap::new(),
+            },
             TimetableImportPathStep {
                 location: TimetableImportPathLocation::OperationalPointLocation { uic: 1 },
                 schedule: HashMap::new(),
@@ -519,5 +604,11 @@ mod tests {
         assert_eq!(ops_uic.len(), 2);
         assert_eq!(ops_uic[0], 1);
         assert_eq!(ops_uic[1], 2);
+
+        let ops_id = ops_id_from_path(&path);
+
+        assert_eq!(ops_id.len(), 2);
+        assert_eq!(ops_id[0], "a1");
+        assert_eq!(ops_id[1], "a2");
     }
 }