Add source support

e2e-tests/src/main/java/gov/nasa/jpl/aerie/e2e/procedural/scheduling/procedures/ExternalEventsSourceQueryGoal.java

@@ -0,0 +1,32 @@
+package gov.nasa.jpl.aerie.e2e.procedural.scheduling.procedures;
+
+import gov.nasa.ammos.aerie.procedural.scheduling.Goal;
+import gov.nasa.ammos.aerie.procedural.scheduling.annotations.SchedulingProcedure;
+import gov.nasa.ammos.aerie.procedural.scheduling.plan.EditablePlan;
+import gov.nasa.ammos.aerie.procedural.timeline.payloads.ExternalSource;
+import gov.nasa.ammos.aerie.procedural.timeline.payloads.activities.DirectiveStart;
+import gov.nasa.ammos.aerie.procedural.timeline.plan.EventQuery;
+import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;
+import java.util.Map;
+
+@SchedulingProcedure
+public record ExternalEventsSourceQueryGoal() implements Goal {
+  @Override
+  public void run(@NotNull final EditablePlan plan) {
+
+    // demonstrate more complicated query functionality
+    EventQuery eventQuery = new EventQuery(
+        null,
+        null,
+        List.of(new ExternalSource("NewTest.json", "TestGroup_2"))
+    );
+
+    for (final var e: plan.events(eventQuery)) {
+      plan.create("BiteBanana", new DirectiveStart.Absolute(e.getInterval().start), Map.of("biteSize", SerializedValue.of(1)));
+    }
+    plan.commit();
+  }
+}

e2e-tests/src/main/java/gov/nasa/jpl/aerie/e2e/procedural/scheduling/procedures/ExternalEventsQueryGoal.java → e2e-tests/src/main/java/gov/nasa/jpl/aerie/e2e/procedural/scheduling/procedures/ExternalEventsTypeQueryGoal.java

@@ -12,7 +12,7 @@
 import java.util.Map;
 
 @SchedulingProcedure
-public record ExternalEventsQueryGoal() implements Goal {
+public record ExternalEventsTypeQueryGoal() implements Goal {
   @Override
   public void run(@NotNull final EditablePlan plan) {
 

e2e-tests/src/test/java/gov/nasa/jpl/aerie/e2e/procedural/scheduling/ExternalEventsTests.java

@@ -164,10 +164,10 @@ void testExternalEventSimple() throws IOException {
   }
 
   @Test
-  void testExternalEventQuery() throws IOException {
+  void testExternalEventTypeQuery() throws IOException {
     // first, run the goal
     try (final var gateway = new GatewayRequests(playwright)) {
-      int procedureJarId = gateway.uploadJarFile("build/libs/ExternalEventsQueryGoal.jar");
+      int procedureJarId = gateway.uploadJarFile("build/libs/ExternalEventsTypeQueryGoal.jar");
       // Add Scheduling Procedure
       procedureId = hasura.createSchedulingSpecProcedure(
           "Test Scheduling Procedure",
@@ -205,4 +205,35 @@ void testExternalEventQuery() throws IOException {
       assertEquals(activityStartTime.toString(), expected.get(i).start_time());
     }
   }
+
+  @Test
+  void testExternalEventSourceQuery() throws IOException {
+    // first, run the goal
+    try (final var gateway = new GatewayRequests(playwright)) {
+      int procedureJarId = gateway.uploadJarFile("build/libs/ExternalEventsSourceQueryGoal.jar");
+      // Add Scheduling Procedure
+      procedureId = hasura.createSchedulingSpecProcedure(
+          "Test Scheduling Procedure",
+          procedureJarId,
+          specId,
+          0
+      );
+    }
+    hasura.awaitScheduling(specId);
+    final var plan = hasura.getPlan(planId);
+    final var activities = plan.activityDirectives();
+
+    // ensure the orderings line up
+    activities.sort(Comparator.comparing(Plan.ActivityDirective::startOffset));
+
+    // compare arrays
+    assertEquals(additionalExternalEvents.size(), activities.size());
+    for (int i = 0; i < activities.size(); i++) {
+      Instant activityStartTime = Duration.addToInstant(
+          Instant.parse(planStartTimestamp),
+          Duration.fromString(activities.get(i).startOffset())
+      );
+      assertEquals(activityStartTime.toString(), additionalExternalEvents.get(i).start_time());
+    }
+  }
 }

procedural/timeline/src/main/kotlin/gov/nasa/ammos/aerie/procedural/timeline/collections/ExternalEvents.kt

@@ -6,6 +6,7 @@ import gov.nasa.ammos.aerie.procedural.timeline.payloads.activities.Instance
 import gov.nasa.ammos.aerie.procedural.timeline.ops.*
 import gov.nasa.ammos.aerie.procedural.timeline.ops.coalesce.CoalesceNoOp
 import gov.nasa.ammos.aerie.procedural.timeline.payloads.ExternalEvent
+import gov.nasa.ammos.aerie.procedural.timeline.payloads.ExternalSource
 import gov.nasa.ammos.aerie.procedural.timeline.util.preprocessList
 
 /**
@@ -23,8 +24,8 @@ data class ExternalEvents(private val timeline: Timeline<ExternalEvent, External
   fun filterByType(vararg types: String) = filter { it.type in types }
 
   /** Filter by one or more event sources. */
-  fun filterBySource(vararg sources: String) = filter { it.source in sources }
+  fun filterBySource(vararg sources: ExternalSource) = filter { it.source in sources }
 
   /** Filter by one or more derivation groups. */
-  fun filterByDerivationGroup(vararg groups: String) = filter { it.derivationGroup in groups }
+  fun filterByDerivationGroup(vararg groups: String) = filter { it.source.derivationGroup in groups }
 }

procedural/timeline/src/main/kotlin/gov/nasa/ammos/aerie/procedural/timeline/payloads/ExternalEvent.kt

@@ -6,14 +6,15 @@ import gov.nasa.jpl.aerie.merlin.protocol.types.SerializedValue
 /** An external event instance. */
 data class ExternalEvent(
   /** The string name of this event. */
+  @JvmField
   val key: String,
   /** The type of the event. */
+  @JvmField
   val type: String,
   /** The source this event comes from. */
-  val source: String,
-  /** The derivation group that this event comes from. */
-  val derivationGroup: String,
+  @JvmField
+  val source: ExternalSource,
   override val interval: Interval,
 ): IntervalLike<ExternalEvent> {
-  override fun withNewInterval(i: Interval) = ExternalEvent(key, type, source, derivationGroup, i)
+  override fun withNewInterval(i: Interval) = ExternalEvent(key, type, source, i)
 }

procedural/timeline/src/main/kotlin/gov/nasa/ammos/aerie/procedural/timeline/payloads/ExternalSource.kt

@@ -0,0 +1,14 @@
+package gov.nasa.ammos.aerie.procedural.timeline.payloads
+
+/**
+ * An external source instance. Used for querying purposes - see EventQuery.kt.
+ * The included fields represent the primary key used to identify External Sources.
+ */
+data class ExternalSource(
+  /** The string name of this source. */
+  @JvmField
+  val key: String,
+  /** The derivation group that this source is a member of. */
+  @JvmField
+  val derivationGroup: String,
+)

procedural/timeline/src/main/kotlin/gov/nasa/ammos/aerie/procedural/timeline/plan/EventQuery.kt

@@ -1,5 +1,7 @@
 package gov.nasa.ammos.aerie.procedural.timeline.plan
 
+import gov.nasa.ammos.aerie.procedural.timeline.payloads.ExternalSource
+
 /** Fields for filtering events as they are queried. */
 data class EventQuery(
   /**
@@ -17,13 +19,14 @@ data class EventQuery(
   val eventTypes: List<String>?,
 
   /**
-   * A nullable list of sources; the event must belong to one of them if present.
+   * A nullable list of sources (described as a tuple of the source's (key, derivation group name)); the event must
+   *    belong to one of them if present.
    *
    * If null, all sources are allowed.
    */
-  val sources: List<String>?,
+  val sources: List<ExternalSource>?,
 ) {
-  constructor(derivationGroup: String?, eventType: String?, source: String?): this(
+  constructor(derivationGroup: String?, eventType: String?, source: ExternalSource?): this(
     derivationGroup?.let { listOf(it) },
     eventType?.let { listOf(it) },
     source?.let { listOf(it) }

scheduler-driver/src/main/java/gov/nasa/jpl/aerie/scheduler/solver/PrioritySolver.java

@@ -46,7 +46,6 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.MICROSECOND;
 import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.ZERO;
-import static gov.nasa.jpl.aerie.merlin.protocol.types.Duration.min;
 
 /**
  * prototype scheduling algorithm that schedules activities for a plan

scheduler-server/src/main/java/gov/nasa/jpl/aerie/scheduler/server/services/GraphQLMerlinDatabaseService.java

@@ -2,6 +2,7 @@
 
 import gov.nasa.ammos.aerie.procedural.timeline.Interval;
 import gov.nasa.ammos.aerie.procedural.timeline.payloads.ExternalEvent;
+import gov.nasa.ammos.aerie.procedural.timeline.payloads.ExternalSource;
 import gov.nasa.jpl.aerie.constraints.model.DiscreteProfile;
 import gov.nasa.jpl.aerie.constraints.model.LinearProfile;
 import gov.nasa.jpl.aerie.json.BasicParsers;
@@ -1044,7 +1045,7 @@ public Map<String, List<ExternalEvent>> getExternalEvents(final PlanId planId, f
     final var unorganized =  parseExternalEvents(data, horizonStart);
     final var result = new HashMap<String, List<ExternalEvent>>();
     for (final var event: unorganized) {
-      final var list = result.computeIfAbsent(event.getDerivationGroup(), $ -> new ArrayList<>());
+      final var list = result.computeIfAbsent(event.source.derivationGroup, $ -> new ArrayList<>());
       list.add(event);
     }
     return result;
@@ -1185,8 +1186,10 @@ private List<ExternalEvent> parseExternalEvents(final JsonArray eventsJson, fina
       result.add(new ExternalEvent(
           e.getString("event_key"),
           e.getString("event_type_name"),
-          e.getString("source_key"),
-          e.getString("derivation_group_name"),
+          new ExternalSource(
+              e.getString("source_key"),
+              e.getString("derivation_group_name")
+          ),
           Interval.between(start, end)
       ));
     }