Implement Temporal Subset Simulation

db-tests/src/test/java/gov/nasa/jpl/aerie/database/MerlinDatabaseTests.java

@@ -130,47 +130,42 @@ INSERT INTO simulation_template (model_id, description, arguments)
     }
   }
 
-  int insertSimulationWithTemplateId(final int simulationTemplateId, final int planId) throws SQLException {
+  int getSimulationId(final int planId) throws SQLException {
     try (final var statement = connection.createStatement()) {
-      final var res = statement
-          .executeQuery(
+      final var res = statement.executeQuery(
               """
-                  INSERT INTO simulation (simulation_template_id, plan_id, arguments)
-                  VALUES ('%s', '%s', '{}')
-                  RETURNING id;"""
-                  .formatted(simulationTemplateId, planId)
+                  SELECT id
+                  FROM simulation
+                  WHERE simulation.plan_id = '%s';
+              """.formatted(planId)
           );
       res.next();
       return res.getInt("id");
     }
   }
 
-  int insertSimulationWithoutTemplateId(final int planId) throws SQLException {
+   void addTemplateIdToSimulation(final int simulationTemplateId, final int simulationId) throws SQLException {
     try (final var statement = connection.createStatement()) {
-      final var res = statement
-          .executeQuery(
-              """
-                  INSERT INTO simulation (plan_id, arguments)
-                  VALUES ('%s', '{}')
-                  RETURNING id;"""
-                  .formatted(planId)
-          );
-      res.next();
-      return res.getInt("id");
+      statement.executeUpdate(
+          """
+              UPDATE simulation
+              SET simulation_template_id = '%s',
+                  arguments = '{}'
+              WHERE id = '%s';
+          """.formatted(simulationTemplateId, simulationId));
     }
   }
 
-  int insertDataset() throws SQLException {
+  int getDatasetId(final int planId) throws SQLException {
     try (final var statement = connection.createStatement()) {
       final var res = statement
           .executeQuery(
               """
-                  INSERT INTO dataset
-                  DEFAULT VALUES
-                  RETURNING id;"""
+                  SELECT dataset_id from plan_dataset
+                  WHERE plan_id = '%s'""".formatted(planId)
           );
       res.next();
-      return res.getInt("id");
+      return res.getInt("dataset_id");
     }
   }
 
@@ -194,8 +189,8 @@ SimulationDatasetRecord insertSimulationDataset(final int simulationId, final in
       final var res = statement
           .executeQuery(
               """
-                  INSERT INTO simulation_dataset (simulation_id, dataset_id, offset_from_plan_start)
-                  VALUES ('%s', '%s', '0')
+                  INSERT INTO simulation_dataset (simulation_id, dataset_id, arguments, simulation_start_time, simulation_end_time)
+                  VALUES ('%s', '%s', '{}', '2020-1-1 00:00:00', '2020-1-2 00:00:00')
                   RETURNING simulation_id, dataset_id;"""
                   .formatted(simulationId, datasetId)
           );
@@ -209,8 +204,7 @@ INSERT INTO simulation_dataset (simulation_id, dataset_id, offset_from_plan_star
   int planId;
   int activityId;
   int simulationTemplateId;
-  int simulationWithTemplateId;
-  int simulationWithoutTemplateId;
+  int simulationId;
   int datasetId;
   SimulationDatasetRecord simulationDatasetRecord;
   PlanDatasetRecord planDatasetRecord;
@@ -222,11 +216,11 @@ void beforeEach() throws SQLException {
     planId = insertPlan(missionModelId);
     activityId = insertActivity(planId);
     simulationTemplateId = insertSimulationTemplate(missionModelId);
-    simulationWithTemplateId = insertSimulationWithTemplateId(simulationTemplateId, planId);
-    simulationWithoutTemplateId = insertSimulationWithoutTemplateId(planId);
+    simulationId = getSimulationId(planId);
+    addTemplateIdToSimulation(simulationTemplateId, simulationId);
     planDatasetRecord = insertPlanDataset(planId);
-    datasetId = insertDataset();
-    simulationDatasetRecord = insertSimulationDataset(simulationWithTemplateId, datasetId);
+    datasetId = getDatasetId(planId);
+    simulationDatasetRecord = insertSimulationDataset(simulationId, datasetId);
   }
 
   @AfterEach
@@ -513,7 +507,7 @@ void shouldIncrementSimulationRevisionOnSimulationUpdate() throws SQLException {
                                            """
                                                SELECT revision FROM simulation
                                                WHERE id = %s;"""
-                                               .formatted(simulationWithTemplateId)
+                                               .formatted(simulationId)
                                        );
       initialRes.next();
       final var initialRevision = initialRes.getInt("revision");
@@ -524,15 +518,15 @@ void shouldIncrementSimulationRevisionOnSimulationUpdate() throws SQLException {
                     """
                         UPDATE simulation SET arguments = '{}'
                         WHERE id = %s;"""
-                        .formatted(simulationWithTemplateId)
+                        .formatted(simulationId)
                 );
 
       final var updatedRes = connection.createStatement()
                                        .executeQuery(
                                            """
                                                SELECT revision FROM simulation
                                                WHERE id = %s;"""
-                                               .formatted(simulationWithTemplateId)
+                                               .formatted(simulationId)
                                        );
       updatedRes.next();
       final var updatedRevision = updatedRes.getInt("revision");
@@ -665,7 +659,7 @@ SELECT COUNT(*) FROM dataset
   @Nested
   class SimulationDatasetTriggers {
     @Test
-    void shouldInitializeDatasetOnInsertWithTemplate() throws SQLException {
+    void shouldInitializeDatasetOnInsert() throws SQLException {
       try (final var statement = connection.createStatement()) {
         try (final var res = statement.executeQuery(
             """
@@ -677,33 +671,12 @@ void shouldInitializeDatasetOnInsertWithTemplate() throws SQLException {
           res.next();
           assertEquals(1, res.getInt("plan_revision"));
           assertEquals(0, res.getInt("model_revision"));
-          assertEquals(0, res.getInt("simulation_revision"));
+          assertEquals(1, res.getInt("simulation_revision")); //1, as we add a template in the BeforeEach
           assertEquals(0, res.getInt("simulation_template_revision"));
         }
       }
     }
 
-    @Test
-    void shouldInitializeDatasetOnInsertWithoutTemplate() throws SQLException {
-      try (final var statement = connection.createStatement()) {
-        final var simulationDatasetWithoutTemplateId = insertSimulationDataset(simulationWithoutTemplateId, datasetId);
-        try (final var res = statement.executeQuery(
-            """
-                SELECT plan_revision, model_revision, simulation_revision, simulation_template_revision
-                FROM simulation_dataset
-                WHERE simulation_id = %s AND dataset_id = %s;"""
-                .formatted(simulationDatasetWithoutTemplateId.simulation_id(), simulationDatasetWithoutTemplateId.dataset_id())
-        )) {
-          res.next();
-          assertEquals(1, res.getInt("plan_revision"));
-          assertEquals(0, res.getInt("model_revision"));
-          assertEquals(0, res.getInt("simulation_revision"));
-          res.getInt("simulation_template_revision");
-          assertTrue(res.wasNull());
-        }
-      }
-    }
-
     @Test
     void shouldDeleteDatasetOnSimulationDatasetDelete() throws SQLException {
       try (final var statement = connection.createStatement()) {
@@ -818,7 +791,7 @@ void shouldCancelSimulationOnSimulationUpdate() throws SQLException {
                     UPDATE simulation
                     SET arguments = '{}'
                     WHERE id = %s;"""
-                    .formatted(simulationWithTemplateId)
+                    .formatted(simulationId)
             );
 
         try (final var res = statement.executeQuery(

deployment/hasura/migrations/AerieMerlin/5_sim_start_end/down.sql

@@ -0,0 +1,9 @@
+alter table simulation
+  drop column simulation_start_time,
+  drop column simulation_end_time;
+
+alter table simulation_template
+  drop column simulation_start_time,
+  drop column simulation_end_time;
+
+call migrations.mark_migration_rolled_back('5');

deployment/hasura/migrations/AerieMerlin/5_sim_start_end/up.sql

@@ -0,0 +1,9 @@
+alter table simulation
+  add column simulation_start_time timestamptz default null,
+  add column simulation_end_time timestamptz default null;
+
+alter table simulation_template
+  add column simulation_start_time timestamptz default null,
+  add column simulation_end_time timestamptz default null;
+
+call migrations.mark_migration_applied('5');

deployment/hasura/migrations/AerieMerlin/6_create_simulation_on_plan_insert/down.sql

@@ -0,0 +1,52 @@
+-- Revert Duplicate Plan
+create or replace function duplicate_plan(plan_id integer, new_plan_name text)
+  returns integer -- plan_id of the new plan
+  security definer
+  language plpgsql as $$
+  declare
+    validate_plan_id integer;
+    new_plan_id integer;
+    created_snapshot_id integer;
+begin
+  select id from plan where plan.id = duplicate_plan.plan_id into validate_plan_id;
+  if(validate_plan_id is null) then
+    raise exception 'Plan % does not exist.', plan_id;
+  end if;
+
+  select create_snapshot(plan_id) into created_snapshot_id;
+
+  insert into plan(revision, name, model_id, duration, start_time, parent_id)
+    select
+        0, new_plan_name, model_id, duration, start_time, plan_id
+    from plan where id = plan_id
+    returning id into new_plan_id;
+  insert into activity_directive(
+      id, plan_id, name, tags, source_scheduling_goal_id, created_at, last_modified_at, start_offset, type, arguments,
+      last_modified_arguments_at, metadata, anchor_id, anchored_to_start
+    )
+    select
+      id, new_plan_id, name, tags, source_scheduling_goal_id, created_at, last_modified_at, start_offset, type, arguments,
+      last_modified_arguments_at, metadata, anchor_id, anchored_to_start
+    from activity_directive where activity_directive.plan_id = duplicate_plan.plan_id;
+  insert into simulation (revision, simulation_template_id, plan_id, arguments)
+    select 0, simulation_template_id, new_plan_id, arguments
+    from simulation
+    where simulation.plan_id = duplicate_plan.plan_id;
+
+  insert into preset_to_directive(preset_id, activity_id, plan_id)
+  select preset_id, activity_id, new_plan_id
+  from preset_to_directive ptd where ptd.plan_id = duplicate_plan.plan_id;
+
+  insert into plan_latest_snapshot(plan_id, snapshot_id) values(new_plan_id, created_snapshot_id);
+  return new_plan_id;
+end;
+$$;
+
+-- Remove uniqueness constraint
+alter table simulation drop constraint one_simulation_per_plan;
+
+-- Remove on insert trigger
+drop trigger simulation_row_for_new_plan_trigger on plan;
+drop function create_simulation_row_for_new_plan();
+
+call migrations.mark_migration_rolled_back('6');

deployment/hasura/migrations/AerieMerlin/6_create_simulation_on_plan_insert/up.sql

@@ -0,0 +1,74 @@
+-- Add a trigger for creating simulation on plan insert
+create function create_simulation_row_for_new_plan()
+returns trigger
+security definer
+language plpgsql as $$begin
+  insert into simulation (revision, simulation_template_id, plan_id, arguments)
+  values (0, null, new.id, '{}');
+  return new;
+end
+$$;
+
+create trigger simulation_row_for_new_plan_trigger
+after insert on plan
+for each row
+execute function create_simulation_row_for_new_plan();
+
+-- Allow only one simulation row per plan
+-- If your migration is failing, this is probably why.
+alter table simulation add constraint one_simulation_per_plan unique(plan_id);
+
+-- Modify Duplicate Plan to update rather than insert
+create or replace function duplicate_plan(plan_id integer, new_plan_name text)
+  returns integer -- plan_id of the new plan
+  security definer
+  language plpgsql as $$
+  declare
+    validate_plan_id integer;
+    new_plan_id integer;
+    created_snapshot_id integer;
+begin
+  select id from plan where plan.id = duplicate_plan.plan_id into validate_plan_id;
+  if(validate_plan_id is null) then
+    raise exception 'Plan % does not exist.', plan_id;
+  end if;
+
+  select create_snapshot(plan_id) into created_snapshot_id;
+
+  insert into plan(revision, name, model_id, duration, start_time, parent_id)
+    select
+        0, new_plan_name, model_id, duration, start_time, plan_id
+    from plan where id = plan_id
+    returning id into new_plan_id;
+  insert into activity_directive(
+      id, plan_id, name, tags, source_scheduling_goal_id, created_at, last_modified_at, start_offset, type, arguments,
+      last_modified_arguments_at, metadata, anchor_id, anchored_to_start
+    )
+    select
+      id, new_plan_id, name, tags, source_scheduling_goal_id, created_at, last_modified_at, start_offset, type, arguments,
+      last_modified_arguments_at, metadata, anchor_id, anchored_to_start
+    from activity_directive where activity_directive.plan_id = duplicate_plan.plan_id;
+
+  with source_plan as (
+    select simulation_template_id, arguments, simulation_start_time, simulation_end_time
+    from simulation
+    where simulation.plan_id = duplicate_plan.plan_id
+  )
+  update simulation s
+  set simulation_template_id = source_plan.simulation_template_id,
+      arguments = source_plan.arguments,
+      simulation_start_time = source_plan.simulation_start_time,
+      simulation_end_time = source_plan.simulation_end_time
+  from source_plan
+  where s.plan_id = new_plan_id;
+
+  insert into preset_to_directive(preset_id, activity_id, plan_id)
+  select preset_id, activity_id, new_plan_id
+  from preset_to_directive ptd where ptd.plan_id = duplicate_plan.plan_id;
+
+  insert into plan_latest_snapshot(plan_id, snapshot_id) values(new_plan_id, created_snapshot_id);
+  return new_plan_id;
+end;
+$$;
+
+call migrations.mark_migration_applied('6');

deployment/hasura/migrations/AerieMerlin/7_store_sim_args_in_dataset/down.sql

@@ -0,0 +1,53 @@
+-- Reset the View
+drop view simulated_activity;
+create view simulated_activity as
+(
+  select span.id as id,
+         simulation_dataset.id as simulation_dataset_id,
+         span.parent_id as parent_id,
+         span.start_offset as start_offset,
+         span.duration as duration,
+         span.attributes as attributes,
+         span.type as activity_type_name,
+         (span.attributes#>>'{directiveId}')::integer as directive_id,
+         plan.start_time + span.start_offset as start_time,
+         plan.start_time + span.start_offset + span.duration as end_time
+
+   from span
+     join dataset on span.dataset_id = dataset.id
+     join simulation_dataset on dataset.id = simulation_dataset.dataset_id
+     join simulation on simulation.id = simulation_dataset.simulation_id
+     join plan on plan.id = simulation.plan_id
+);
+comment on view simulated_activity is e''
+  'Concrete activity instance created via simulation.';
+comment on column simulated_activity.id is e''
+  'Unique identifier for the activity instance span.';
+comment on column simulated_activity.simulation_dataset_id is e''
+  'The simulation dataset this activity is part of.';
+comment on column simulated_activity.parent_id is e''
+  'The parent activity of this activity.';
+comment on column simulated_activity.start_offset is e''
+  'The offset from the dataset start at which this activity begins.';
+comment on column simulated_activity.duration is e''
+  'The amount of time this activity extends for.';
+comment on column simulated_activity.attributes is e''
+  'A set of named values annotating this activity.';
+comment on column simulated_activity.activity_type_name is e''
+  'The activity type of this activity.';
+comment on column simulated_activity.directive_id is e''
+  'The id of the activity directive that created this activity.';
+comment on column simulated_activity.start_time is e''
+  'The absolute start time of this activity.';
+comment on column simulated_activity.end_time is e''
+  'The absolute end time of this activity.';
+
+-- Drop the trigger
+drop trigger update_offset_from_plan_start_trigger on simulation_dataset;
+drop function update_offset_from_plan_start();
+
+-- Drop the new columns
+alter table simulation_dataset
+drop column arguments,
+drop column simulation_start_time,
+drop column simulation_end_time;