AWS SDK for Java で S3 のファイルをまとめてダウンロードしたい

S3 のフォルダにあるファイルを zip でまとめてダウンロードしたい。今回は AWS SDK for JavaAmazonS3#getObject を非同期で呼び出して ZipOutputStream で書き出してみる。

サンプルコードはこちら。

github.com

前提

備忘録

起動クラスに@EnableAsync アノテーションを付けて非同期処理を有効にする。

@EnableAsync
public class DemoApplication {
    public static void main(String[] args) {
        SpringApplication.run(DemoApplication.class, args);
    }
    // ...(略)

任意の prefix に一致するオブジェクトリストを取得する。サブフォルダは除く。

public List<String> listObjectKeys(String bucketName, String prefix) {
    return amazonS3.listObjectsV2(bucketName, prefix)
            .getObjectSummaries().stream()
            .filter(object -> !object.getKey().endsWith("/")) // サブフォルダは除く
            .map(object -> object.getKey())
            .collect(Collectors.toList());
}

AmazonS3#getObject を非同期で呼び出す。戻り値には CompletableFuture<T> を返す。ドキュメントを読む限りでは AmazonS3 はスレッドセーフっぽい。(AmazonS3ClientBuilder の方はスレッドアンセーフ)

サービスクライアントの作成 - AWS SDK for Java

@Async
public CompletableFuture<S3Content> fetchAsync(String bucketName, String objectKey) {
    S3Object object = amazonS3.getObject(new GetObjectRequest(bucketName, objectKey));
    return CompletableFuture.completedFuture(new S3Content(object));
}

オブジェクトリストの objectKey ごとに @Async な非同期処理を呼び出して処理完了を待つ。

List<S3Content> s3Contents = new ArrayList<>();
List<CompletableFuture<S3Content>> processes = new ArrayList<>();
for (String objectKey : objectKeys) {
    CompletableFuture<S3Content> process =
            service.fetchAsync(bucketName, objectKey).whenCompleteAsync((res, e) -> s3Contents.add(res));
    processes.add(process);
}
CompletableFuture.allOf(processes.toArray(new CompletableFuture[objectKeys.size()])).join();

byte 配列のデータを ZipOutputStream で書き出す。今回は ServletResponse#getOutputStream をラップして呼び出し元に返すようにする。

String basename = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyyMMdd-HHmmss"));

response.setContentType("application/octet-stream");
response.setHeader("Content-Disposition", "attachment;filename=" + basename + ".zip");
response.setStatus(HttpServletResponse.SC_OK);

try (ZipOutputStream out = new ZipOutputStream(response.getOutputStream())) {
    for (S3Content content : s3Contents) {
        out.putNextEntry(new ZipEntry(Paths.get(basename, content.filename()).toString()));
        try (InputStream in = new ByteArrayInputStream(content.data())) {
            byte[] buf = new byte[4096];
            int len;
            while ((len = in.read(buf)) > 0) {
                out.write(buf, 0, len);
            }
        }
    }
}