import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.model.ListObjectsV2Request;
import com.amazonaws.services.s3.model.S3ObjectSummary;
import lombok.val;

import java.util.*;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

class S3ObjectSummaryIterator implements Iterator<S3ObjectSummary> {

    private final AmazonS3 client;
    private final ListObjectsV2Request request;
    private final Queue<S3ObjectSummary> summaries = new LinkedList<>();
    private boolean isTruncated = false;

    S3ObjectSummaryIterator(AmazonS3 client, ListObjectsV2Request request) {
        this.client = client;
        this.request = request;
        fetch();
    }

    private void fetch() {
        val result = this.client.listObjectsV2(this.request);
        summaries.addAll(result.getObjectSummaries());
        this.isTruncated = result.isTruncated();
        if (result.isTruncated()) {
            request.setContinuationToken(result.getNextContinuationToken());
        }
    }

    @Override
    public boolean hasNext() {
        if (summaries.isEmpty() && isTruncated) {
            fetch();
        }
        return !summaries.isEmpty();
    }

    @Override
    public S3ObjectSummary next() {
        return summaries.poll();
    }

    public static Stream<S3ObjectSummary> stream(AmazonS3 client, ListObjectsV2Request request) {
        val iterator = new S3ObjectSummaryIterator(client, request);
        return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false);
    }

}

Recommended Posts