[DC-1117] Fix dataset name validation for createSnapshot (#1703)

* Fix dataset name validation for createSnapshot

* Fetch datasets only once

* Make source dataset a parameter for createSnapshot flight

* Only pass around ID

* Store ID

* Pull lists

* Fix compilation

* Remove old comment

* Fix spotless

* Fix tests

* Add bad request exception and content validation

* Fix merge

* Fix merge again and remove unused code

* Fix grammar

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

@@ -6,6 +6,7 @@
 import bio.terra.app.utils.ControllerUtils;
 import bio.terra.common.SqlSortDirection;
 import bio.terra.common.ValidationUtils;
+import bio.terra.common.exception.BadRequestException;
 import bio.terra.common.iam.AuthenticatedUserRequest;
 import bio.terra.common.iam.AuthenticatedUserRequestFactory;
 import bio.terra.controller.SnapshotsApi;
@@ -39,8 +40,8 @@
 import bio.terra.service.auth.iam.IamAction;
 import bio.terra.service.auth.iam.IamResourceType;
 import bio.terra.service.auth.iam.IamService;
-import bio.terra.service.auth.iam.exception.IamForbiddenException;
 import bio.terra.service.dataset.AssetModelValidator;
+import bio.terra.service.dataset.Dataset;
 import bio.terra.service.dataset.IngestRequestValidator;
 import bio.terra.service.filedata.FileService;
 import bio.terra.service.job.JobService;
@@ -122,32 +123,21 @@ private AuthenticatedUserRequest getAuthenticatedInfo() {
     return authenticatedUserRequestFactory.from(request);
   }
 
-  // -- snapshot --
-  private List<UUID> getUnauthorizedSources(
-      List<UUID> snapshotSourceDatasetIds, AuthenticatedUserRequest userReq) {
-    return snapshotSourceDatasetIds.stream()
-        .filter(
-            sourceId ->
-                !iamService.isAuthorized(
-                    userReq, IamResourceType.DATASET, sourceId.toString(), IamAction.LINK_SNAPSHOT))
-        .collect(Collectors.toList());
-  }
-
   @Override
   public ResponseEntity<JobModel> createSnapshot(
       @Valid @RequestBody SnapshotRequestModel snapshotRequestModel) {
     AuthenticatedUserRequest userReq = getAuthenticatedInfo();
-    List<UUID> snapshotSourceDatasetIds =
-        snapshotService.getSourceDatasetIdsFromSnapshotRequest(snapshotRequestModel);
-    // TODO: auth should be put into flight
-    List<UUID> unauthorized = getUnauthorizedSources(snapshotSourceDatasetIds, userReq);
-    if (unauthorized.isEmpty()) {
-      String jobId = snapshotService.createSnapshot(snapshotRequestModel, userReq);
-      // we can retrieve the job we just created
-      return jobToResponse(jobService.retrieveJob(jobId, userReq));
+
+    if (snapshotRequestModel.getContents().size() > 1) {
+      throw new BadRequestException("Snapshot request must contain at most one source.");
     }
-    throw new IamForbiddenException(
-        "User is not authorized to create snapshots for these datasets " + unauthorized);
+
+    Dataset dataset = snapshotService.getSourceDatasetFromSnapshotRequest(snapshotRequestModel);
+    iamService.verifyAuthorization(
+        userReq, IamResourceType.DATASET, dataset.getId().toString(), IamAction.LINK_SNAPSHOT);
+    String jobId = snapshotService.createSnapshot(snapshotRequestModel, dataset, userReq);
+    // we can retrieve the job we just created
+    return jobToResponse(jobService.retrieveJob(jobId, userReq));
   }
 
   @Override

src/main/java/bio/terra/service/resourcemanagement/ResourceService.java

@@ -513,32 +513,21 @@ public void updateBucketMetadata(
   public UUID initializeSnapshotProject(
       BillingProfileModel billingProfile,
       String projectId,
-      List<Dataset> sourceDatasets,
+      Dataset sourceDataset,
       String snapshotName,
       UUID snapshotId)
       throws InterruptedException {
 
     GoogleRegion region =
         (GoogleRegion)
-            sourceDatasets
-                .iterator()
-                .next()
+            sourceDataset
                 .getDatasetSummary()
                 .getStorageResourceRegion(GoogleCloudResource.FIRESTORE);
 
-    String datasetNames =
-        sourceDatasets.stream().map(Dataset::getName).collect(Collectors.joining(","));
-
-    String datasetIds =
-        sourceDatasets.stream()
-            .map(Dataset::getId)
-            .map(UUID::toString)
-            .collect(Collectors.joining(","));
-
     Map<String, String> labels =
         Map.of(
-            "dataset-names", datasetNames,
-            "dataset-ids", datasetIds,
+            "dataset-names", sourceDataset.getName(),
+            "dataset-ids", sourceDataset.getId().toString(),
             "snapshot-name", snapshotName,
             "snapshot-id", snapshotId.toString(),
             "project-usage", "snapshot");

src/main/java/bio/terra/service/snapshot/SnapshotService.java

@@ -174,10 +174,9 @@ public SnapshotService(
    * @return jobId (flightId) of the job
    */
   public String createSnapshot(
-      SnapshotRequestModel snapshotRequestModel, AuthenticatedUserRequest userReq) {
-    SnapshotRequestContentsModel requestContents = snapshotRequestModel.getContents().get(0);
-    String sourceDatasetName = requestContents.getDatasetName();
-    Dataset dataset = datasetService.retrieveByName(sourceDatasetName);
+      SnapshotRequestModel snapshotRequestModel,
+      Dataset dataset,
+      AuthenticatedUserRequest userReq) {
     if (snapshotRequestModel.getProfileId() == null) {
       snapshotRequestModel.setProfileId(dataset.getDefaultProfileId());
       logger.warn(
@@ -200,9 +199,7 @@ public String createSnapshot(
     return jobService
         .newJob(description, SnapshotCreateFlight.class, snapshotRequestModel, userReq)
         .addParameter(CommonMapKeys.CREATED_AT, Instant.now().toEpochMilli())
-        .addParameter(JobMapKeys.IAM_RESOURCE_TYPE.getKeyName(), IamResourceType.DATASET)
-        .addParameter(JobMapKeys.IAM_RESOURCE_ID.getKeyName(), dataset.getId())
-        .addParameter(JobMapKeys.IAM_ACTION.getKeyName(), IamAction.LINK_SNAPSHOT)
+        .addParameter(JobMapKeys.DATASET_ID.getKeyName(), dataset.getId())
         .addParameter(JobMapKeys.SNAPSHOT_ID.getKeyName(), snapshotId.toString())
         .submit();
   }
@@ -653,26 +650,15 @@ public static AssetSpecification getAssetByNameFromDataset(Dataset dataset, Stri
                     "This dataset does not have an asset specification with name: " + assetName));
   }
 
-  public List<UUID> getSourceDatasetIdsFromSnapshotRequest(
-      SnapshotRequestModel snapshotRequestModel) {
-    return getSourceDatasetsFromSnapshotRequest(snapshotRequestModel).stream()
-        .map(Dataset::getId)
-        .collect(Collectors.toList());
-  }
-
-  public List<Dataset> getSourceDatasetsFromSnapshotRequest(
-      SnapshotRequestModel snapshotRequestModel) {
-    return snapshotRequestModel.getContents().stream()
-        .map(
-            c ->
-                c.getMode() == SnapshotRequestContentsModel.ModeEnum.BYREQUESTID
-                    ? retrieve(
-                            snapshotRequestDao
-                                .getById(c.getRequestIdSpec().getSnapshotRequestId())
-                                .getSourceSnapshotId())
-                        .getSourceDataset()
-                    : datasetService.retrieveByName(c.getDatasetName()))
-        .collect(Collectors.toList());
+  public Dataset getSourceDatasetFromSnapshotRequest(SnapshotRequestModel snapshotRequestModel) {
+    SnapshotRequestContentsModel contents = snapshotRequestModel.getContents().get(0);
+    return contents.getMode() == SnapshotRequestContentsModel.ModeEnum.BYREQUESTID
+        ? retrieve(
+                snapshotRequestDao
+                    .getById(contents.getRequestIdSpec().getSnapshotRequestId())
+                    .getSourceSnapshotId())
+            .getSourceDataset()
+        : datasetService.retrieveByName(contents.getDatasetName());
   }
 
   public AddAuthDomainResponseModel addSnapshotDataAccessControls(

src/main/java/bio/terra/service/snapshot/flight/create/CreateSnapshotInitializeProjectStep.java

@@ -11,22 +11,21 @@
 import bio.terra.stairway.Step;
 import bio.terra.stairway.StepResult;
 import bio.terra.stairway.StepStatus;
-import java.util.List;
 import java.util.UUID;
 
 public class CreateSnapshotInitializeProjectStep implements Step {
   private final ResourceService resourceService;
-  private final List<Dataset> sourceDatasets;
+  private final Dataset sourceDataset;
   private final String snapshotName;
   private final UUID snapshotId;
 
   public CreateSnapshotInitializeProjectStep(
       ResourceService resourceService,
-      List<Dataset> sourceDatasets,
+      Dataset sourceDataset,
       String snapshotName,
       UUID snapshotId) {
     this.resourceService = resourceService;
-    this.sourceDatasets = sourceDatasets;
+    this.sourceDataset = sourceDataset;
     this.snapshotName = snapshotName;
     this.snapshotId = snapshotId;
   }
@@ -44,7 +43,7 @@ public StepResult doStep(FlightContext context) throws InterruptedException {
     try {
       projectResourceId =
           resourceService.initializeSnapshotProject(
-              profileModel, projectId, sourceDatasets, snapshotName, snapshotId);
+              profileModel, projectId, sourceDataset, snapshotName, snapshotId);
     } catch (GoogleResourceException e) {
       if (e.getCause().getMessage().contains("500 Internal Server Error")) {
         return new StepResult(StepStatus.STEP_RESULT_FAILURE_RETRY, e);

src/main/java/bio/terra/service/snapshot/flight/create/SnapshotCreateFlight.java

@@ -134,11 +134,8 @@ public SnapshotCreateFlight(FlightMap inputParameters, Object applicationContext
     RetryRule randomBackoffRetry =
         getDefaultRandomBackoffRetryRule(appConfig.getMaxStairwayThreads());
 
-    // TODO note that with multi-dataset snapshots this will need to change
-    List<Dataset> sourceDatasets =
-        snapshotService.getSourceDatasetsFromSnapshotRequest(snapshotReq);
-    Dataset sourceDataset = sourceDatasets.get(0);
-    UUID datasetId = sourceDataset.getId();
+    UUID datasetId = inputParameters.get(JobMapKeys.DATASET_ID.getKeyName(), UUID.class);
+    Dataset sourceDataset = datasetService.retrieve(datasetId);
     String datasetName = sourceDataset.getName();
 
     var platform =
@@ -166,7 +163,7 @@ public SnapshotCreateFlight(FlightMap inputParameters, Object applicationContext
       // Get or initialize the project where the snapshot resources will be created
       addStep(
           new CreateSnapshotInitializeProjectStep(
-              resourceService, sourceDatasets, snapshotName, snapshotId),
+              resourceService, sourceDataset, snapshotName, snapshotId),
           getDefaultExponentialBackoffRetryRule());
     }
 
@@ -208,7 +205,7 @@ public SnapshotCreateFlight(FlightMap inputParameters, Object applicationContext
     }
 
     // Make the big query dataset with views and populate row id filtering tables.
-    // Depending on the type of snapshot, the primary data step will differ:
+    // Depending on the type of snapshot, the primary data step will diff
     // TODO: this assumes single-dataset snapshots, will need to add a loop for multiple
     switch (mode) {
       case BYASSET -> {

src/test/java/bio/terra/app/controller/SnapshotsApiControllerTest.java

@@ -50,6 +50,7 @@
 import bio.terra.service.auth.iam.IamService;
 import bio.terra.service.auth.iam.exception.IamForbiddenException;
 import bio.terra.service.dataset.AssetModelValidator;
+import bio.terra.service.dataset.Dataset;
 import bio.terra.service.dataset.IngestRequestValidator;
 import bio.terra.service.filedata.FileService;
 import bio.terra.service.job.JobService;
@@ -301,16 +302,14 @@ void unlockSnapshot() throws Exception {
   @Test
   void createSnapshot() throws Exception {
     mockValidators();
-
-    when(snapshotService.getSourceDatasetIdsFromSnapshotRequest(SNAPSHOT_REQUEST_MODEL))
-        .thenReturn(List.of(DATASET_ID));
+    Dataset dataset = new Dataset().id(DATASET_ID);
+    when(snapshotService.getSourceDatasetFromSnapshotRequest(SNAPSHOT_REQUEST_MODEL))
+        .thenReturn(dataset);
 
     IamAction iamAction = IamAction.LINK_SNAPSHOT;
-    when(iamService.isAuthorized(
-            TEST_USER, IamResourceType.DATASET, DATASET_ID.toString(), iamAction))
-        .thenReturn(true);
 
-    when(snapshotService.createSnapshot(SNAPSHOT_REQUEST_MODEL, TEST_USER)).thenReturn(JOB_ID);
+    when(snapshotService.createSnapshot(SNAPSHOT_REQUEST_MODEL, dataset, TEST_USER))
+        .thenReturn(JOB_ID);
     when(jobService.retrieveJob(JOB_ID, TEST_USER)).thenReturn(JOB_MODEL);
 
     String actualJson =
@@ -323,32 +322,28 @@ void createSnapshot() throws Exception {
             .getResponse()
             .getContentAsString();
     JobModel actual = TestUtils.mapFromJson(actualJson, JobModel.class);
-    assertThat("Job model is returned", actual, equalTo(JOB_MODEL));
-
     verify(iamService)
-        .isAuthorized(TEST_USER, IamResourceType.DATASET, DATASET_ID.toString(), iamAction);
+        .verifyAuthorization(TEST_USER, IamResourceType.DATASET, DATASET_ID.toString(), iamAction);
+    assertThat("Job model is returned", actual, equalTo(JOB_MODEL));
   }
 
   @Test
   void createSnapshot_forbidden() throws Exception {
     mockValidators();
-
-    when(snapshotService.getSourceDatasetIdsFromSnapshotRequest(SNAPSHOT_REQUEST_MODEL))
-        .thenReturn(List.of(DATASET_ID));
+    Dataset dataset = new Dataset().id(DATASET_ID);
+    when(snapshotService.getSourceDatasetFromSnapshotRequest(SNAPSHOT_REQUEST_MODEL))
+        .thenReturn(dataset);
 
     IamAction iamAction = IamAction.LINK_SNAPSHOT;
-    when(iamService.isAuthorized(
-            TEST_USER, IamResourceType.DATASET, DATASET_ID.toString(), iamAction))
-        .thenReturn(false);
+    doThrow(IamForbiddenException.class)
+        .when(iamService)
+        .verifyAuthorization(TEST_USER, IamResourceType.DATASET, DATASET_ID.toString(), iamAction);
 
     mvc.perform(
             post(SNAPSHOTS_ENDPOINT)
                 .contentType(MediaType.APPLICATION_JSON)
                 .content(TestUtils.mapToJson(SNAPSHOT_REQUEST_MODEL)))
         .andExpect(status().isForbidden());
-
-    verify(iamService)
-        .isAuthorized(TEST_USER, IamResourceType.DATASET, DATASET_ID.toString(), iamAction);
   }
 
   @Test

src/test/java/bio/terra/service/snapshot/SnapshotRequestValidatorTest.java

@@ -15,10 +15,12 @@
 import bio.terra.model.JobModel;
 import bio.terra.service.auth.iam.IamService;
 import bio.terra.service.dataset.AssetModelValidator;
+import bio.terra.service.dataset.Dataset;
 import bio.terra.service.dataset.IngestRequestValidator;
 import bio.terra.service.filedata.FileService;
 import bio.terra.service.job.JobService;
 import bio.terra.service.snapshotbuilder.SnapshotBuilderService;
+import java.util.UUID;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
@@ -86,6 +88,8 @@ void validateByRequestIdFail() throws Exception {
   void validateByRequestId() throws Exception {
     when(jobService.retrieveJob(any(), any()))
         .thenReturn(new JobModel().jobStatus(JobModel.JobStatusEnum.SUCCEEDED));
+    when(snapshotService.getSourceDatasetFromSnapshotRequest(any()))
+        .thenReturn(new Dataset().id(UUID.randomUUID()));
     String validRequestModel =
         """
         {