Add SimDatasetId override to lookup()

merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/exceptions/SimulationDatasetMismatchException.java

@@ -0,0 +1,15 @@
+package gov.nasa.jpl.aerie.merlin.server.exceptions;
+
+import gov.nasa.jpl.aerie.merlin.server.models.PlanId;
+import gov.nasa.jpl.aerie.merlin.server.models.SimulationDatasetId;
+
+public class SimulationDatasetMismatchException extends Exception {
+  public final SimulationDatasetId simDatasetId;
+  public final PlanId planId;
+
+  public SimulationDatasetMismatchException(final PlanId pid, final SimulationDatasetId sid) {
+    super("Simulation Dataset with id `" + sid.id() + "` does not belong to Plan with id `"+pid.id()+"`");
+    this.planId = pid;
+    this.simDatasetId = sid;
+  }
+}

merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/http/ResponseSerializers.java

@@ -16,6 +16,7 @@
 import gov.nasa.jpl.aerie.merlin.protocol.types.ValueSchema;
 import gov.nasa.jpl.aerie.merlin.server.exceptions.NoSuchPlanDatasetException;
 import gov.nasa.jpl.aerie.merlin.server.exceptions.NoSuchPlanException;
+import gov.nasa.jpl.aerie.merlin.server.exceptions.SimulationDatasetMismatchException;
 import gov.nasa.jpl.aerie.merlin.server.remotes.MissionModelAccessException;
 import gov.nasa.jpl.aerie.merlin.server.services.GetSimulationResultsAction;
 import gov.nasa.jpl.aerie.merlin.server.services.LocalMissionModelService;
@@ -468,6 +469,13 @@ public static JsonValue serializeInputMismatchException(final InputMismatchExcep
                .build();
   }
 
+  public static JsonValue serializeSimulationDatasetMismatchException(final SimulationDatasetMismatchException ex){
+     return Json.createObjectBuilder()
+               .add("message", "input mismatch exception")
+               .add("cause", ex.getMessage())
+               .build();
+  }
+
   private static final class ValueSchemaSerializer implements ValueSchema.Visitor<JsonValue> {
     @Override
     public JsonValue onReal() {

merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/InMemoryResultsCellRepository.java

@@ -61,6 +61,11 @@ public Optional<ResultsProtocol.ReaderRole> lookup(final PlanId planId) {
     }
   }
 
+  @Override
+  public Optional<ResultsProtocol.ReaderRole> lookup(final PlanId planId, final SimulationDatasetId simulationDatasetId) {
+    return lookup(planId);
+  }
+
   public boolean isEqualTo(final InMemoryResultsCellRepository other) {
     return this.cells.equals(other.cells);
   }

merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/ResultsCellRepository.java

@@ -1,7 +1,9 @@
 package gov.nasa.jpl.aerie.merlin.server.remotes;
 
 import gov.nasa.jpl.aerie.merlin.server.ResultsProtocol;
+import gov.nasa.jpl.aerie.merlin.server.exceptions.SimulationDatasetMismatchException;
 import gov.nasa.jpl.aerie.merlin.server.models.PlanId;
+import gov.nasa.jpl.aerie.merlin.server.models.SimulationDatasetId;
 
 import java.util.Optional;
 
@@ -11,4 +13,6 @@ public interface ResultsCellRepository {
   Optional<ResultsProtocol.OwnerRole> claim(PlanId planId, Long datasetId);
 
   Optional<ResultsProtocol.ReaderRole> lookup(PlanId planId);
+
+  Optional<ResultsProtocol.ReaderRole> lookup(PlanId planId, SimulationDatasetId simulationDatasetId) throws SimulationDatasetMismatchException;
 }

merlin-server/src/main/java/gov/nasa/jpl/aerie/merlin/server/remotes/postgres/PostgresResultsCellRepository.java

@@ -12,6 +12,7 @@
 import gov.nasa.jpl.aerie.merlin.protocol.types.ValueSchema;
 import gov.nasa.jpl.aerie.merlin.server.ResultsProtocol;
 import gov.nasa.jpl.aerie.merlin.server.ResultsProtocol.State;
+import gov.nasa.jpl.aerie.merlin.server.exceptions.SimulationDatasetMismatchException;
 import gov.nasa.jpl.aerie.merlin.server.models.PlanId;
 import gov.nasa.jpl.aerie.merlin.server.models.ProfileSet;
 import gov.nasa.jpl.aerie.merlin.server.models.SimulationDatasetId;
@@ -120,14 +121,32 @@ public Optional<ResultsProtocol.OwnerRole> claim(final PlanId planId, final Long
   public Optional<ResultsProtocol.ReaderRole> lookup(final PlanId planId) {
     try (final var connection = this.dataSource.getConnection()) {
       final var simulation = getSimulation(connection, planId);
-      final var datasetRecord = lookupSimulationDatasetRecord(
-          connection,
-          simulation.id());
-      final var datasetId$ = datasetRecord.map(SimulationDatasetRecord::datasetId);
+      final var datasetRecord = lookupSimulationDatasetRecord(connection, simulation.id());
+
+      if (datasetRecord.isEmpty()) return Optional.empty();
+
+      final var datasetId = datasetRecord.get().datasetId();
+      return Optional.of(new PostgresResultsCell(this.dataSource, simulation, datasetId));
+    } catch (final SQLException ex) {
+      throw new DatabaseException("Failed to get simulation", ex);
+    }
+  }
 
-      if (datasetId$.isEmpty()) return Optional.empty();
+  @Override
+  public Optional<ResultsProtocol.ReaderRole> lookup(final PlanId planId, final SimulationDatasetId simulationDatasetId) throws SimulationDatasetMismatchException {
+    try (final var connection = this.dataSource.getConnection()) {
+      final var simulation = getSimulation(connection, planId);
+      final var datasetRecord = getSimulationDatasetRecordById(connection, simulationDatasetId.id());
+
+      if (datasetRecord.isEmpty()) return Optional.empty();
+      // If this check fails, then the specified sim dataset is not a simulation for the specified plan
+      if(datasetRecord.get().simulationId() != simulation.id()){
+        throw new SimulationDatasetMismatchException(
+            planId,
+            new SimulationDatasetId(datasetRecord.get().simulationDatasetId()));
+      }
 
-      final var datasetId = datasetId$.get();
+      final var datasetId = datasetRecord.get().datasetId();
       return Optional.of(new PostgresResultsCell(this.dataSource, simulation, datasetId));
     } catch (final SQLException ex) {
       throw new DatabaseException("Failed to get simulation", ex);