Yet another approach to type safe step parameters

src/main/java/bio/terra/app/controller/DatasetsApiController.java

@@ -345,7 +345,7 @@ public ResponseEntity<JobModel> addDatasetAssetSpecifications(
       @PathVariable("id") UUID id, @Valid @RequestBody AssetModel asset) {
     AuthenticatedUserRequest userReq = getAuthenticatedInfo();
     verifyDatasetAuthorization(userReq, id.toString(), IamAction.MANAGE_SCHEMA);
-    String jobId = datasetService.addDatasetAssetSpecifications(id.toString(), asset, userReq);
+    String jobId = datasetService.addDatasetAssetSpecifications(id, asset, userReq);
     return jobToResponse(jobService.retrieveJob(jobId, userReq));
   }
 

src/main/java/bio/terra/common/BaseStep.java

@@ -0,0 +1,75 @@
+package bio.terra.common;
+
+import bio.terra.model.ErrorModel;
+import bio.terra.stairway.FlightContext;
+import bio.terra.stairway.Step;
+import bio.terra.stairway.StepResult;
+import bio.terra.stairway.exception.RetryException;
+import com.google.common.annotations.VisibleForTesting;
+import org.springframework.http.HttpStatus;
+
+public abstract class BaseStep implements Step {
+  @StepOutput protected Object response;
+  @StepOutput protected HttpStatus statusCode;
+
+  private FlightContext context;
+
+  @Override
+  public final StepResult doStep(FlightContext context)
+      throws InterruptedException, RetryException {
+    this.context = context;
+    StepUtils.readInputs(this, context);
+    try {
+      return perform();
+    } finally {
+      StepUtils.writeOutputs(this, context);
+    }
+  }
+
+  @Override
+  public final StepResult undoStep(FlightContext context) throws InterruptedException {
+    this.context = context;
+    StepUtils.readInputs(this, context);
+    return undo();
+  }
+
+  public abstract StepResult perform() throws InterruptedException, RetryException;
+
+  public StepResult undo() throws InterruptedException {
+    // Many steps aren't undoable.
+    return StepResult.getStepResultSuccess();
+  }
+
+  protected void setErrorResponse(String message, HttpStatus responseStatus) {
+    ErrorModel errorModel = new ErrorModel().message(message);
+    setResponse(errorModel, responseStatus);
+  }
+
+  protected void setResponse(Object responseObject, HttpStatus responseStatus) {
+    response = responseObject;
+    statusCode = responseStatus;
+  }
+
+  protected void setResponse(Object responseObject) {
+    response = responseObject;
+    statusCode = HttpStatus.OK;
+  }
+
+  protected FlightContext getContext() {
+    return context;
+  }
+
+  protected String getFlightId() {
+    return context.getFlightId();
+  }
+
+  @VisibleForTesting
+  public Object getResponse() {
+    return response;
+  }
+
+  @VisibleForTesting
+  public HttpStatus getStatusCode() {
+    return statusCode;
+  }
+}

src/main/java/bio/terra/common/FlightUtils.java

@@ -1,7 +1,5 @@
 package bio.terra.common;
 
-import bio.terra.model.ErrorModel;
-import bio.terra.service.job.JobMapKeys;
 import bio.terra.stairway.FlightContext;
 import bio.terra.stairway.FlightMap;
 import bio.terra.stairway.RetryRuleExponentialBackoff;
@@ -11,41 +9,13 @@
 import org.apache.commons.lang3.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.springframework.http.HttpStatus;
 
 /** Common methods for building flights */
 public final class FlightUtils {
   private static final Logger logger = LoggerFactory.getLogger(FlightUtils.class);
 
   private FlightUtils() {}
 
-  /**
-   * Build an error model and set it as the response
-   *
-   * @param context
-   * @param message
-   * @param responseStatus
-   */
-  public static void setErrorResponse(
-      FlightContext context, String message, HttpStatus responseStatus) {
-    ErrorModel errorModel = new ErrorModel().message(message);
-    setResponse(context, errorModel, responseStatus);
-  }
-
-  /**
-   * Set the response and status code in the result map.
-   *
-   * @param context flight context
-   * @param responseObject response object to set
-   * @param responseStatus status code to set
-   */
-  public static void setResponse(
-      FlightContext context, Object responseObject, HttpStatus responseStatus) {
-    FlightMap workingMap = context.getWorkingMap();
-    workingMap.put(JobMapKeys.RESPONSE.getKeyName(), responseObject);
-    workingMap.put(JobMapKeys.STATUS_CODE.getKeyName(), responseStatus);
-  }
-
   /**
    * Common logic for deciding if a BigQuery exception is a retry-able IAM propagation error. There
    * is not a specific reason code for the IAM setPolicy failed error. This check is a bit fragile.

src/main/java/bio/terra/common/StepInput.java

@@ -0,0 +1,14 @@
+package bio.terra.common;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.FIELD)
+public @interface StepInput {
+  String USE_DEFAULT_NAME = "";
+
+  String value() default USE_DEFAULT_NAME;
+}

src/main/java/bio/terra/common/StepOutput.java

@@ -0,0 +1,14 @@
+package bio.terra.common;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.FIELD)
+public @interface StepOutput {
+  String USE_DEFAULT_NAME = "";
+
+  String value() default USE_DEFAULT_NAME;
+}

src/main/java/bio/terra/common/StepUtils.java

@@ -0,0 +1,92 @@
+package bio.terra.common;
+
+import bio.terra.stairway.FlightContext;
+import bio.terra.stairway.FlightMap;
+import bio.terra.stairway.Step;
+import java.lang.reflect.Field;
+
+// Suppress sonar warnings about reflection API usage. Reflection APIs must be used to read and
+// write fields in a Step.
+@SuppressWarnings({"java:S3011"})
+public class StepUtils {
+
+  public static class MissingStepInputException extends RuntimeException {
+    public MissingStepInputException(String key) {
+      super("No flight value found for StepInput key '" + key + "'");
+    }
+  }
+
+  public static class IllegalSetException extends RuntimeException {
+    public IllegalSetException(Throwable cause) {
+      super(cause);
+    }
+  }
+
+  public static class IllegalGetException extends RuntimeException {
+    public IllegalGetException(Throwable cause) {
+      super(cause);
+    }
+  }
+
+  private StepUtils() {}
+
+  public static String keyFromField(Field field) {
+    var input = field.getAnnotation(StepInput.class);
+    if (input != null && !input.value().isEmpty()) {
+      return input.value();
+    }
+    var output = field.getAnnotation(StepOutput.class);
+    if (output != null && !output.value().isEmpty()) {
+      return output.value();
+    }
+    return field.getName();
+  }
+
+  public static void readInputs(Step step, FlightContext context) throws MissingStepInputException {
+    for (Class<?> clazz = step.getClass(); clazz != null; clazz = clazz.getSuperclass()) {
+      for (Field field : clazz.getDeclaredFields()) {
+        if (field.isAnnotationPresent(StepInput.class)) {
+          String key = keyFromField(field);
+          if (context.getInputParameters().containsKey(key)) {
+            setField(step, context.getInputParameters(), field, key);
+          } else if (context.getWorkingMap().containsKey(key)) {
+            setField(step, context.getWorkingMap(), field, key);
+          } else if (!field.isAnnotationPresent(StepOutput.class)) {
+            // If the field is only used as an input, report an error if there's no value for it.
+            throw new MissingStepInputException(key);
+          }
+        }
+      }
+    }
+  }
+
+  private static void setField(Step step, FlightMap map, Field field, String key) {
+    field.setAccessible(true);
+    try {
+      field.set(step, map.get(key, field.getType()));
+    } catch (IllegalAccessException e) {
+      throw new IllegalSetException(e);
+    }
+  }
+
+  public static void writeOutputs(Step step, FlightContext context) {
+    for (Class<?> clazz = step.getClass(); clazz != null; clazz = clazz.getSuperclass()) {
+      for (Field field : clazz.getDeclaredFields()) {
+        if (field.isAnnotationPresent(StepOutput.class)) {
+          field.setAccessible(true);
+          final Object value;
+          try {
+            value = field.get(step);
+          } catch (IllegalAccessException e) {
+            throw new IllegalGetException(e);
+          }
+          if (value == null) {
+            // An unset output can occur if an exception is thrown inside the run() operation.
+            continue;
+          }
+          context.getWorkingMap().put(keyFromField(field), value);
+        }
+      }
+    }
+  }
+}

src/main/java/bio/terra/service/common/CommonFlightUtils.java

@@ -22,6 +22,13 @@ public static Instant getCreatedAt(FlightContext context) {
     return null;
   }
 
+  public static Instant getCreatedAt(Long createdAtRaw) {
+    if (createdAtRaw != null) {
+      return Instant.ofEpochMilli(createdAtRaw);
+    }
+    return null;
+  }
+
   /**
    * Obtains flight information of Interest that might be useful in a journal entry. Patterned after
    * an allowList because the flightContext input parameters can contain some sensitive information

src/main/java/bio/terra/service/common/ResourceLockSetResponseStep.java

@@ -1,20 +1,16 @@
 package bio.terra.service.common;
 
-import bio.terra.common.FlightUtils;
+import bio.terra.common.BaseStep;
 import bio.terra.model.ResourceLocks;
-import bio.terra.service.job.DefaultUndoStep;
-import bio.terra.stairway.FlightContext;
 import bio.terra.stairway.StepResult;
-import org.springframework.http.HttpStatus;
 
-public abstract class ResourceLockSetResponseStep extends DefaultUndoStep {
+public abstract class ResourceLockSetResponseStep extends BaseStep {
 
   protected ResourceLockSetResponseStep() {}
 
   @Override
-  public StepResult doStep(FlightContext context) {
-    ResourceLocks locks = getResourceLocks();
-    FlightUtils.setResponse(context, locks, HttpStatus.OK);
+  public StepResult perform() {
+    setResponse(getResourceLocks());
     return StepResult.getStepResultSuccess();
   }
 

src/main/java/bio/terra/service/common/gcs/BigQueryUtils.java

@@ -6,7 +6,11 @@
 public class BigQueryUtils {
 
   public static String getSuffix(FlightContext context) {
-    return context.getFlightId().replace('-', '_');
+    return getSuffix(context.getFlightId());
+  }
+
+  public static String getSuffix(String flightId) {
+    return flightId.replace('-', '_');
   }
 
   public static String gsPathMappingTableName(Snapshot snapshot) {

src/main/java/bio/terra/service/dataset/DatasetService.java

@@ -418,7 +418,7 @@ public String ingestDataset(
   }
 
   public String addDatasetAssetSpecifications(
-      String datasetId, AssetModel assetModel, AuthenticatedUserRequest userReq) {
+      UUID datasetId, AssetModel assetModel, AuthenticatedUserRequest userReq) {
     String description = "Add dataset asset specification " + assetModel.getName();
     return jobService
         .newJob(description, AddAssetSpecFlight.class, assetModel, userReq)