Skip to content

Commit

Permalink
feat: enable download to not run in parallel
Browse files Browse the repository at this point in the history
  • Loading branch information
jfroche committed Apr 21, 2022
1 parent 693f49f commit e718db4
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 6 deletions.
21 changes: 18 additions & 3 deletions src/main/java/de/taimos/pipeline/aws/S3DownloadStep.java
Original file line number Diff line number Diff line change
Expand Up @@ -52,12 +52,14 @@ public class S3DownloadStep extends AbstractS3Step {
private final String bucket;
private String path = "";
private boolean force = false;
private boolean disableParallelDownloads = false;

@DataBoundConstructor
public S3DownloadStep(String file, String bucket, boolean pathStyleAccessEnabled, boolean payloadSigningEnabled) {
public S3DownloadStep(String file, String bucket, boolean pathStyleAccessEnabled, boolean payloadSigningEnabled, boolean disableParallelDownloads) {
super(pathStyleAccessEnabled, payloadSigningEnabled);
this.file = file;
this.bucket = bucket;
this.disableParallelDownloads = disableParallelDownloads;
}

public String getFile() {
Expand All @@ -76,6 +78,10 @@ public boolean isForce() {
return this.force;
}

public boolean isDisableParallelDownloads() {
return this.disableParallelDownloads;
}

@DataBoundSetter
public void setForce(boolean force) {
this.force = force;
Expand All @@ -86,6 +92,11 @@ public void setPath(String path) {
this.path = path;
}

@DataBoundSetter
public void setDisableParallelDownload(boolean disableParallelDownloads) {
this.disableParallelDownloads = disableParallelDownloads;
}

@Override
public StepExecution start(StepContext context) throws Exception {
return new S3DownloadStep.Execution(this, context);
Expand Down Expand Up @@ -130,6 +141,7 @@ public Void run() throws Exception {
final String bucket = this.step.getBucket();
final String path = this.step.getPath();
final boolean force = this.step.isForce();
final boolean disableParallelDownloads = this.step.isDisableParallelDownloads();

Preconditions.checkArgument(bucket != null && !bucket.isEmpty(), "Bucket must not be null or empty");

Expand All @@ -146,7 +158,7 @@ public Void run() throws Exception {
throw new RuntimeException("Target exists: " + target.toURI().toString());
}
}
target.act(new RemoteDownloader(Execution.this.step.createS3ClientOptions(), envVars, listener, bucket, path));
target.act(new RemoteDownloader(Execution.this.step.createS3ClientOptions(), envVars, listener, bucket, path, disableParallelDownloads));
listener.getLogger().println("Download complete");
return null;
}
Expand All @@ -162,19 +174,22 @@ private static class RemoteDownloader extends MasterToSlaveFileCallable<Void> {
private final TaskListener taskListener;
private final String bucket;
private final String path;
private final Boolean disableParallelDownloads;

RemoteDownloader(S3ClientOptions amazonS3ClientOptions, EnvVars envVars, TaskListener taskListener, String bucket, String path) {
RemoteDownloader(S3ClientOptions amazonS3ClientOptions, EnvVars envVars, TaskListener taskListener, String bucket, String path, Boolean disableParallelDownloads) {
this.amazonS3ClientOptions = amazonS3ClientOptions;
this.envVars = envVars;
this.taskListener = taskListener;
this.bucket = bucket;
this.path = path;
this.disableParallelDownloads = disableParallelDownloads;
}

@Override
public Void invoke(File localFile, VirtualChannel channel) throws IOException, InterruptedException {
TransferManager mgr = TransferManagerBuilder.standard()
.withS3Client(AWSClientFactory.create(this.amazonS3ClientOptions.createAmazonS3ClientBuilder(), this.envVars))
.withDisableParallelDownloads(this.disableParallelDownloads)
.build();

if (this.path == null || this.path.isEmpty() || this.path.endsWith("/")) {
Expand Down
6 changes: 3 additions & 3 deletions src/test/java/de/taimos/pipeline/aws/S3DownloadStepTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -27,20 +27,20 @@
public class S3DownloadStepTest {
@Test
public void gettersWorkAsExpected() throws Exception {
S3DownloadStep step = new S3DownloadStep("my-file", "my-bucket", false, false);
S3DownloadStep step = new S3DownloadStep("my-file", "my-bucket", false, false, false);
Assert.assertEquals("my-file", step.getFile());
Assert.assertEquals("my-bucket", step.getBucket());
}

@Test
public void defaultPathIsEmpty() throws Exception {
S3DownloadStep step = new S3DownloadStep("my-file", "my-bucket", false, false);
S3DownloadStep step = new S3DownloadStep("my-file", "my-bucket", false, false, false);
Assert.assertEquals("", step.getPath());
}

@Test
public void defaultForceIsFalse() throws Exception {
S3DownloadStep step = new S3DownloadStep("my-file", "my-bucket", false, false);
S3DownloadStep step = new S3DownloadStep("my-file", "my-bucket", false, false, false);
Assert.assertFalse(step.isForce());
}
}

0 comments on commit e718db4

Please sign in to comment.