Merge pull request #782 from HubSpot/downloader-custom-creds

support per-bucket creds for downloading artifacts
HubSpot · Dec 29, 2015 · 6625612 · 6625612
2 parents 9634c33 + 93ce48c
commit 6625612
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/S3ArtifactDownloader.java b/SingularityS3Base/src/main/java/com/hubspot/singularity/s3/base/S3ArtifactDownloader.java
@@ -57,13 +57,21 @@ public void download(S3Artifact s3Artifact, Path downloadTo) {
     }
   }
 
+  private AWSCredentials getCredentialsForBucket(String bucketName) {
+    if (configuration.getS3BucketCredentials().containsKey(bucketName)) {
+      return configuration.getS3BucketCredentials().get(bucketName).toAWSCredentials();
+    }
+
+    return new AWSCredentials(configuration.getS3AccessKey().get(), configuration.getS3SecretKey().get());
+  }
+
   private void downloadThrows(final S3Artifact s3Artifact, final Path downloadTo) throws Exception {
     log.info("Downloading {}", s3Artifact);
 
     Jets3tProperties jets3tProperties = Jets3tProperties.getInstance(Constants.JETS3T_PROPERTIES_FILENAME);
     jets3tProperties.setProperty("httpclient.socket-timeout-ms", Long.toString(configuration.getS3ChunkDownloadTimeoutMillis()));
 
-    final S3Service s3 = new RestS3Service(new AWSCredentials(configuration.getS3AccessKey().get(), configuration.getS3SecretKey().get()), null, null, jets3tProperties);
+    final S3Service s3 = new RestS3Service(getCredentialsForBucket(s3Artifact.getS3Bucket()), null, null, jets3tProperties);
 
     long length = 0;
 

diff --git a/...Base/src/main/java/com/hubspot/singularity/s3/base/config/SingularityS3Configuration.java b/...Base/src/main/java/com/hubspot/singularity/s3/base/config/SingularityS3Configuration.java
@@ -2,6 +2,8 @@
 
 import static com.hubspot.mesos.JavaUtils.obfuscateValue;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
 import javax.validation.constraints.Min;
@@ -57,6 +59,10 @@ public class SingularityS3Configuration extends BaseRunnerConfiguration {
   @JsonProperty
   private String localDownloadPath = "/download";
 
+  @NotNull
+  @JsonProperty
+  private Map<String, SingularityS3Credentials> s3BucketCredentials = new HashMap<>();
+
   public SingularityS3Configuration() {
     super(Optional.<String>absent());
   }
@@ -133,6 +139,14 @@ public void setS3ChunkDownloadTimeoutMillis(long s3ChunkDownloadTimeoutMillis) {
     this.s3ChunkDownloadTimeoutMillis = s3ChunkDownloadTimeoutMillis;
   }
 
+  public Map<String, SingularityS3Credentials> getS3BucketCredentials() {
+    return s3BucketCredentials;
+  }
+
+  public void setS3BucketCredentials(Map<String, SingularityS3Credentials> s3BucketCredentials) {
+    this.s3BucketCredentials = s3BucketCredentials;
+  }
+
   @Override
   public String toString() {
     return "SingularityS3Configuration[" +
@@ -145,6 +159,7 @@ public String toString() {
             ", s3ChunkRetries=" + s3ChunkRetries +
             ", localDownloadHttpPort=" + localDownloadHttpPort +
             ", localDownloadPath='" + localDownloadPath + '\'' +
+            ", s3BucketCredentials=" + s3BucketCredentials +
             ']';
   }
 }