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

bloussou · bloussou · commit 2fc00075d7b9 · 2024-01-26T15:09:54.000+01:00
diff --git a/editoast/openapi.yaml b/editoast/openapi.yaml
@@ -3324,13 +3324,18 @@ components:
       - properties:
           OperationalPointNotFound:
             properties:
+              missing_ids:
+                items:
+                  type: string
+                type: array
               missing_uics:
                 items:
                   format: int64
                   type: integer
                 type: array
             required:
             - missing_uics
+            - missing_ids
             type: object
         required:
         - OperationalPointNotFound
@@ -3404,6 +3409,17 @@ components:
         - uic
         - type
         type: object
+      - properties:
+          id:
+            type: string
+          type:
+            enum:
+            - operational_point_id
+            type: string
+        required:
+        - id
+        - type
+        type: object
     TimetableImportPathSchedule:
       properties:
         arrival_time:
diff --git a/editoast/src/models/infra_objects/operational_point.rs b/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_obj_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?)
+    }
 }
diff --git a/editoast/src/views/timetable/import.rs b/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)]
@@ -76,10 +78,19 @@ struct TimetableImportReport {
 
 #[derive(Debug, Clone, Serialize, ToSchema)]
 enum TimetableImportError {
-    RollingStockNotFound { name: String },
-    OperationalPointNotFound { missing_uics: Vec<i64> },
-    PathfindingError { cause: InternalError },
-    SimulationError { cause: InternalError },
+    RollingStockNotFound {
+        name: String,
+    },
+    OperationalPointNotFound {
+        missing_uics: Vec<i64>,
+        missing_ids: Vec<String>,
+    },
+    PathfindingError {
+        cause: InternalError,
+    },
+    SimulationError {
+        cause: InternalError,
+    },
 }
 
 #[derive(Debug, Clone, Deserialize, ToSchema)]
@@ -183,19 +194,21 @@ 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,
-        Err(err) => {
-            return Ok(import_item
-                .trains
-                .into_iter()
-                .map(|train: TimetableImportTrain| (train.name, err.clone()))
-                .collect());
-        }
-    };
+    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
+                    .into_iter()
+                    .map(|train: TimetableImportTrain| (train.name, err.clone()))
+                    .collect());
+            }
+        };
     // 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 +308,63 @@ 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 ops_from_uic: Vec<OperationalPointModel> =
+        OperationalPointModel::retrieve_from_uic(conn, infra_id, ops_uic).await?;
+    let mut op_uic_to_parts = HashMap::<_, Vec<_>>::new();
+    for op in ops_from_uic {
+        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_uic_to_parts.len() != ops_uic.len() {
+        ops_uic.iter().for_each(|uic| {
+            if !op_uic_to_parts.contains_key(uic) {
+                missing_uics.push(*uic)
+            }
+        });
+    }
+
+    let ops_from_ids = OperationalPointModel::retrieve_from_obj_ids(conn, infra_id, ops_id).await?;
     let mut op_id_to_parts = HashMap::<_, Vec<_>>::new();
-    for op in ops {
+    for op in ops_from_ids {
         op_id_to_parts
-            .entry(op.data.0.extensions.identifier.unwrap().uic)
+            .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
-    if op_id_to_parts.len() != ops_uic.len() {
-        let missing_uics = ops_uic
-            .iter()
-            .filter(|uic| !op_id_to_parts.contains_key(uic))
-            .cloned()
-            .collect();
+    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 +381,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 +407,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 +505,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 +522,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 +551,17 @@ 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 +571,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 +595,12 @@ 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 +612,18 @@ 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 +635,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");
     }
 }
diff --git a/front/src/common/api/osrdEditoastApi.ts b/front/src/common/api/osrdEditoastApi.ts
@@ -2305,6 +2305,7 @@ export type TimetableImportError =
     }
   | {
       OperationalPointNotFound: {
+        missing_ids: string[];
         missing_uics: number[];
       };
     }
@@ -2332,6 +2333,10 @@ export type TimetableImportPathLocation =
   | {
       type: 'operational_point';
       uic: number;
+    }
+  | {
+      id: string;
+      type: 'operational_point_id';
     };
 export type TimetableImportPathSchedule = {
   arrival_time: string;
