Merge pull request #676 from chuguoxipo/develop

Initial commit for redrive upload API - implementation & unit tests

Author: DwayneJengSage

src/main/java/org/sagebionetworks/bridge/models/upload/UploadCompletionClient.java

@@ -14,5 +14,6 @@ public enum UploadCompletionClient {
      * Upload has been completed by a worker process that listens for the addition 
      * of upload files to the file upload bucket on S3.
      */
-    S3_WORKER
+    S3_WORKER,
+    REDRIVE
 }

src/main/java/org/sagebionetworks/bridge/models/upload/UploadRedriveList.java

@@ -0,0 +1,20 @@
+package org.sagebionetworks.bridge.models.upload;
+
+import java.util.List;
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
+
+import org.sagebionetworks.bridge.models.BridgeEntity;
+
+public class UploadRedriveList implements BridgeEntity {
+    private final List<String> uploadIds;
+
+    @JsonCreator
+    public UploadRedriveList(@JsonProperty("uploadIds") List<String> uploadIds) {
+        this.uploadIds = uploadIds;
+    }
+    public List<String> getUploadIds() {
+        return (uploadIds == null) ? ImmutableList.of() : uploadIds;
+    }
+}

src/main/java/org/sagebionetworks/bridge/models/worker/UploadRedriveWorkerRequest.java

@@ -0,0 +1,32 @@
+package org.sagebionetworks.bridge.models.worker;
+
+/** Worker request to redrive uploads . */
+public class UploadRedriveWorkerRequest {
+    private String s3Bucket;
+    private String s3Key;
+    private String redriveTypeStr;
+
+    public String getS3Bucket() {
+        return s3Bucket;
+    }
+
+    public String getS3Key() {
+        return s3Key;
+    }
+
+    public String getRedriveTypeStr() {
+        return redriveTypeStr;
+    }
+
+    public void setS3Bucket(String s3Bucket) {
+        this.s3Bucket = s3Bucket;
+    }
+
+    public void setS3Key(String s3Key) {
+        this.s3Key = s3Key;
+    }
+
+    public void setRedriveTypeStr(String redriveTypeStr) {
+        this.redriveTypeStr = redriveTypeStr;
+    }
+}

src/main/java/org/sagebionetworks/bridge/services/UploadService.java

@@ -12,7 +12,9 @@
 import static org.sagebionetworks.bridge.BridgeConstants.API_APP_ID;
 import static org.sagebionetworks.bridge.BridgeConstants.API_DEFAULT_PAGE_SIZE;
 import static org.sagebionetworks.bridge.BridgeConstants.CANNOT_BE_BLANK;
+import static org.sagebionetworks.bridge.BridgeUtils.COMMA_SPACE_JOINER;
 
+import java.io.IOException;
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Date;
@@ -30,14 +32,20 @@
 import com.amazonaws.services.s3.model.AmazonS3Exception;
 import com.amazonaws.services.s3.model.GeneratePresignedUrlRequest;
 import com.amazonaws.services.s3.model.ObjectMetadata;
+import com.amazonaws.services.sqs.AmazonSQS;
+import com.amazonaws.services.sqs.model.SendMessageResult;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.base.Stopwatch;
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableSet;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
+
+import org.sagebionetworks.bridge.BridgeConstants;
+import org.sagebionetworks.bridge.dao.HealthCodeDao;
 import org.sagebionetworks.bridge.exceptions.EntityNotFoundException;
 import org.sagebionetworks.bridge.models.PagedResourceList;
 import org.sagebionetworks.bridge.models.accounts.Account;
@@ -63,7 +71,11 @@
 import org.sagebionetworks.bridge.models.apps.App;
 import org.sagebionetworks.bridge.models.healthdata.HealthDataRecordEx3;
 import org.sagebionetworks.bridge.models.schedules2.timelines.TimelineMetadataView;
+import org.sagebionetworks.bridge.models.upload.UploadRedriveList;
 import org.sagebionetworks.bridge.models.upload.UploadViewEx3;
+import org.sagebionetworks.bridge.models.worker.UploadRedriveWorkerRequest;
+import org.sagebionetworks.bridge.models.worker.WorkerRequest;
+import org.sagebionetworks.bridge.s3.S3Helper;
 import org.sagebionetworks.bridge.time.DateUtils;
 import org.sagebionetworks.bridge.models.ForwardCursorPagedResourceList;
 import org.sagebionetworks.bridge.models.accounts.StudyParticipant;
@@ -88,9 +100,12 @@ public class UploadService {
 
     // package-scoped to be available in unit tests
     static final String CONFIG_KEY_UPLOAD_BUCKET = "upload.bucket";
+    static final String CONFIG_KEY_BACKFILL_BUCKET = "backfill.bucket";
+    static final String REDRIVE_UPLOAD_S3_KEY_PREFIX = "redrive-upload-id-";
     static final String METADATA_KEY_EVENT_TIMESTAMP = "eventTimestamp";
     static final String METADATA_KEY_INSTANCE_GUID = "instanceGuid";
     static final String METADATA_KEY_STARTED_ON = "startedOn";
+    static final String WORKER_NAME_UPLOAD_REDRIVE = "UploadRedriveWorker";
 
     private AccountService accountService;
     private AdherenceService adherenceService;
@@ -103,9 +118,15 @@ public class UploadService {
     private Schedule2Service schedule2Service;
     private StudyService studyService;
     private String uploadBucket;
+    private String redriveUploadBucket;
     private UploadDao uploadDao;
     private UploadDedupeDao uploadDedupeDao;
     private UploadValidationService uploadValidationService;
+    private HealthCodeDao healthCodeDao;
+    private String workerQueueUrl;
+    private AmazonSQS sqsClient;
+    private S3Helper s3Helper;
+    private BridgeConfig config;
 
     // These parameters can be overriden to facilitate testing.
     // By default, we sleep 5 seconds, including right at the start and end. This means on our 7th iteration,
@@ -136,7 +157,9 @@ public final void setExporter3Service(Exporter3Service exporter3Service) {
     /** Sets parameters from the specified Bridge config. */
     @Autowired
     final void setConfig(BridgeConfig config) {
+        this.config = config;
         uploadBucket = config.getProperty(CONFIG_KEY_UPLOAD_BUCKET);
+        redriveUploadBucket = config.getProperty(CONFIG_KEY_BACKFILL_BUCKET);
     }
 
     /**
@@ -189,6 +212,21 @@ final void setUploadValidationService(UploadValidationService uploadValidationSe
         this.uploadValidationService = uploadValidationService;
     }
 
+    @Autowired
+    final void setHealthCodeDao(HealthCodeDao healthCodeDao) {
+        this.healthCodeDao = healthCodeDao;
+    }
+
+    @Autowired
+    final void setSqsClient(AmazonSQS sqsClient) {
+        this.sqsClient = sqsClient;
+    }
+
+    @Resource(name = "s3Helper")
+    public final void setS3Helper(S3Helper s3Helper) {
+        this.s3Helper = s3Helper;
+    }
+
     /**
      * Number of iterations while polling for validation status before we time out. This is used primarily by tests to
      * reduce the amount of wait time during tests.
@@ -527,6 +565,88 @@ public void uploadComplete(String appId, UploadCompletionClient completedBy, Upl
         // Save uploadedOn date and uploadId to related adherence records.
         updateAdherenceWithUploadInfo(appId, upload);
     }
+
+    public void redriveUpload(UploadRedriveList redriveList) throws IOException {
+        checkNotNull(redriveList);
+
+        if (redriveList.getUploadIds().isEmpty()) {
+            throw new BadRequestException("No upload ids submitted for redrive");
+        }
+
+        int redriveSize = redriveList.getUploadIds().size();
+
+        logger.info("Redrive uploads" + " with " + redriveSize + " upload ids");
+
+        // Redrive upload.
+        if (redriveSize <= 10) {
+            redriveSmallAmountOfUploads(redriveList.getUploadIds());
+        } else {
+            redriveLargeAmountOfUploads(redriveList.getUploadIds());
+        }
+    }
+
+    private void redriveSmallAmountOfUploads(List<String> uploadIds) throws JsonProcessingException {
+        for (String uploadId : uploadIds) {
+            Upload upload = getUpload(uploadId);
+            String appId = upload.getAppId();
+            if (appId == null) {
+                appId = healthCodeDao.getAppId(upload.getHealthCode());
+            }
+            uploadComplete(appId, UploadCompletionClient.REDRIVE, upload, true);
+        }
+
+        for (String uploadId: uploadIds) {
+            UploadValidationStatus validationStatus = pollUploadValidationStatusUntilComplete(uploadId);
+            if (validationStatus.getStatus() != UploadStatus.SUCCEEDED) {
+                logErrorMessage(uploadId, validationStatus);
+            }
+        }
+    }
+
+    // For unit test.
+    protected void logErrorMessage(String uploadId, UploadValidationStatus validationStatus) {
+        logger.error("Redrive failed for uploadId=" + uploadId + ": " + COMMA_SPACE_JOINER.join(
+                validationStatus.getMessageList()));
+    }
+
+    private void redriveLargeAmountOfUploads(List<String> uploadIds) throws IOException {
+        // set up s3 Key.
+        String currentTime = String.valueOf(DateUtils.getCurrentDateTime().getMillis());
+        String s3Key = REDRIVE_UPLOAD_S3_KEY_PREFIX + currentTime;
+
+        // Upload file to S3 bucket
+        s3Helper.writeLinesToS3(redriveUploadBucket, s3Key, uploadIds);
+
+        // write Json message to sqs
+        // 1. Create request.
+        UploadRedriveWorkerRequest uploadRedriveWorkerRequest = new UploadRedriveWorkerRequest();
+        uploadRedriveWorkerRequest.setS3Key(s3Key);
+        uploadRedriveWorkerRequest.setS3Bucket(redriveUploadBucket);
+        uploadRedriveWorkerRequest.setRedriveTypeStr("upload_id");
+
+        WorkerRequest workerRequest = new WorkerRequest();
+        workerRequest.setService(WORKER_NAME_UPLOAD_REDRIVE);
+        workerRequest.setBody(uploadRedriveWorkerRequest);
+
+        // Convert request to JSON.
+        ObjectMapper objectMapper = BridgeObjectMapper.get();
+        String requestJson;
+        try {
+            requestJson = objectMapper.writeValueAsString(workerRequest);
+        } catch (JsonProcessingException ex) {
+            // This should never happen, but catch and re-throw for code hygiene.
+            throw new BridgeServiceException("Error creating upload redrive request for S3 file " + s3Key,
+                    ex);
+        }
+
+        // 2. Send to SQS.
+        // Note: SqsInitializer runs after Spring, so we need to grab the queue URL dynamically.
+        workerQueueUrl = config.getProperty(BridgeConstants.CONFIG_KEY_WORKER_SQS_URL);
+
+        SendMessageResult sqsResult = sqsClient.sendMessage(workerQueueUrl, requestJson);
+        logger.info("Sent redrive upload request for file " + s3Key +
+                sqsResult.getMessageId());
+    }
 
     public void deleteUploadsForHealthCode(String healthCode) {
         checkArgument(isNotBlank(healthCode));

src/main/java/org/sagebionetworks/bridge/spring/controllers/UploadController.java

@@ -7,6 +7,7 @@
 import static org.sagebionetworks.bridge.Roles.WORKER;
 import static org.springframework.http.MediaType.APPLICATION_JSON_VALUE;
 
+import java.io.IOException;
 import java.util.EnumSet;
 import java.util.Optional;
 
@@ -20,16 +21,19 @@
 import org.springframework.web.bind.annotation.PostMapping;
 import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.RestController;
+
 import org.sagebionetworks.bridge.Roles;
 import org.sagebionetworks.bridge.dao.HealthCodeDao;
 import org.sagebionetworks.bridge.exceptions.EntityNotFoundException;
 import org.sagebionetworks.bridge.exceptions.UnauthorizedException;
 import org.sagebionetworks.bridge.models.Metrics;
 import org.sagebionetworks.bridge.models.RequestInfo;
+import org.sagebionetworks.bridge.models.StatusMessage;
 import org.sagebionetworks.bridge.models.accounts.UserSession;
 import org.sagebionetworks.bridge.models.healthdata.HealthDataRecord;
 import org.sagebionetworks.bridge.models.upload.Upload;
 import org.sagebionetworks.bridge.models.upload.UploadCompletionClient;
+import org.sagebionetworks.bridge.models.upload.UploadRedriveList;
 import org.sagebionetworks.bridge.models.upload.UploadRequest;
 import org.sagebionetworks.bridge.models.upload.UploadSession;
 import org.sagebionetworks.bridge.models.upload.UploadValidationStatus;
@@ -49,6 +53,8 @@ public class UploadController extends BaseController {
 
     private HealthCodeDao healthCodeDao;
 
+    static final StatusMessage REDRIVE_COMPLETE_MSG = new StatusMessage("Upload redrive completed.");
+
     @Autowired
     final void setUploadService(UploadService uploadService) {
         this.uploadService = uploadService;
@@ -141,7 +147,7 @@ public String uploadComplete(@PathVariable String uploadId,
             if (appId == null) {
                 appId = healthCodeDao.getAppId(upload.getHealthCode());
             }
-            uploadCompletionClient = UploadCompletionClient.S3_WORKER;
+            uploadCompletionClient = redrive? UploadCompletionClient.REDRIVE : UploadCompletionClient.S3_WORKER;
         } else {
             // Or, the consented user that originally made the upload request. Check that health codes match.
             // Do not need to look up the app.
@@ -167,6 +173,16 @@ public String uploadComplete(@PathVariable String uploadId,
         // Upload validation status may contain the health data record. Use the filter to filter out health code.
         return HealthDataRecord.PUBLIC_RECORD_WRITER.writeValueAsString(validationStatus);
     }
+
+    @PostMapping("/v3/uploads/redrive")
+    public StatusMessage redriveUploads() throws IOException {
+        getAuthenticatedSession(Roles.SUPERADMIN);
+
+        UploadRedriveList redriveList = parseJson(UploadRedriveList.class);
+        uploadService.redriveUpload(redriveList);
+
+        return REDRIVE_COMPLETE_MSG;
+    }
 
     @GetMapping("/v3/uploads/{uploadId}")
     public UploadView getUpload(@PathVariable String uploadId) {

src/main/resources/BridgeServer2.conf

@@ -109,6 +109,12 @@ attachment.bucket = org-sagebridge-attachment-${bucket.suffix}
 # Exporter 3 Health Data buckets
 health.data.bucket.raw = org-sagebridge-rawhealthdata-${bucket.suffix}
 
+# Backfill buckets
+local.backfill.bucket = org-sagebridge-backfill-devlocal
+dev.backfill.bucket = org-sagebridge-backfill-devdevelop
+uat.backfill.bucket = org-sagebridge-backfill-devstaging
+prod.backfill.bucket = org-sagebridge-backfill-prod
+
 # Upload CMS certificate information
 upload.cms.certificate.country = US
 upload.cms.certificate.state = WA

src/test/java/org/sagebionetworks/bridge/models/schedules2/adherence/AdherenceRecordListTest.java

@@ -20,6 +20,7 @@ public void canSerialize() throws Exception {
         AdherenceRecordList list = new AdherenceRecordList(ImmutableList.of(rec1, rec2));
 
         JsonNode node = BridgeObjectMapper.get().valueToTree(list);
+
         assertEquals(node.get("records").size(), 2);
         // just verify these are adherence records, which we test separately
         assertEquals(node.get("records").get(0).get("clientTimeZone").textValue(), "America/Los_Angeles");

src/test/java/org/sagebionetworks/bridge/models/upload/UploadRedriveListTest.java

@@ -0,0 +1,37 @@
+package org.sagebionetworks.bridge.models.upload;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.google.common.collect.ImmutableList;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import org.sagebionetworks.bridge.json.BridgeObjectMapper;
+
+public class UploadRedriveListTest {
+    private static final String UPLOAD_ID_1 = "upload1";
+    private static final String UPLOAD_ID_2 = "upload2";
+
+    @Test
+    public void canSerialize() throws Exception {
+        UploadRedriveList list = new UploadRedriveList(ImmutableList.of(UPLOAD_ID_1, UPLOAD_ID_2));
+        
+        JsonNode node = BridgeObjectMapper.get().valueToTree(list);
+
+        assertEquals(node.get("uploadIds").size(), 2);
+        // just verify these are adherence records, which we test separately
+        assertEquals(node.get("uploadIds").get(0).textValue(), UPLOAD_ID_1);
+        assertEquals(node.get("uploadIds").get(1).textValue(), UPLOAD_ID_2);
+        assertEquals(node.get("type").textValue(), "UploadRedriveList");
+
+        UploadRedriveList deser = BridgeObjectMapper.get().readValue(node.toString(), UploadRedriveList.class);
+        assertEquals(deser.getUploadIds().size(), 2);
+        assertEquals(deser.getUploadIds().get(0), UPLOAD_ID_1);
+        assertEquals(deser.getUploadIds().get(1), UPLOAD_ID_2);
+    }
+    
+    @Test
+    public void nullList() {
+        UploadRedriveList list = new UploadRedriveList(null);
+        assertEquals(list.getUploadIds(), ImmutableList.of());
+    }
+}