S3ObjectSummary stream improved

Nikita Sapozhnikov · Nikita Sapozhnikov · commit e02b54b62326 · 2016-04-26T15:33:38.000+01:00
diff --git a/micro-s3/src/main/java/com/aol/micro/server/s3/DirectoryCleaner.java b/micro-s3/src/main/java/com/aol/micro/server/s3/DirectoryCleaner.java
@@ -1,5 +1,6 @@
 package com.aol.micro.server.s3;
 
+import java.io.File;
 import java.io.IOException;
 import java.nio.file.FileSystems;
 import java.nio.file.Files;
@@ -15,18 +16,15 @@
 public class DirectoryCleaner {
 
 	private final String temporaryDirectory;
-	private final boolean cleanupOnStart;
 
 	@Autowired
-	public DirectoryCleaner(@Value("${s3.temp.dir:#{null}") String temporaryDirectory,
-			@Value("${s3.temp.clean.on.start:false") boolean cleanupOnStart) {
+	public DirectoryCleaner(@Value("${s3.temp.dir:#{null}") String temporaryDirectory) {
 		this.temporaryDirectory = temporaryDirectory;
-		this.cleanupOnStart = cleanupOnStart;
 	}
 
 	@PostConstruct
 	public void clean() throws IOException {
-		if (cleanupOnStart) {
+		if (temporaryDirectory != null && new File(temporaryDirectory).exists()) {
 			Path directory = FileSystems.getDefault().getPath(temporaryDirectory);
 			Files.walkFileTree(directory, new CleanupFileVisitor(directory));
 		}
diff --git a/micro-s3/src/main/java/com/aol/micro/server/s3/S3ObjectSummaryIterator.java b/micro-s3/src/main/java/com/aol/micro/server/s3/S3ObjectSummaryIterator.java
@@ -0,0 +1,49 @@
+package com.aol.micro.server.s3;
+
+import java.util.Iterator;
+
+import com.amazonaws.services.s3.AmazonS3Client;
+import com.amazonaws.services.s3.model.ListObjectsRequest;
+import com.amazonaws.services.s3.model.ObjectListing;
+import com.amazonaws.services.s3.model.S3ObjectSummary;
+
+public class S3ObjectSummaryIterator implements Iterator<S3ObjectSummary> {
+
+	private final AmazonS3Client client;
+	private ListObjectsRequest req;
+	private Iterator<S3ObjectSummary> iterator;
+	private boolean empty = true;
+	public S3ObjectSummaryIterator(AmazonS3Client client, ListObjectsRequest req) {
+		this.client = client;
+		this.req = req;
+		updateIterator();
+	}
+	
+	@Override
+	public boolean hasNext() {
+		if(iterator.hasNext()) {
+			return true;
+		} else if(!empty){
+			updateIterator();
+			return iterator.hasNext();
+		} else {
+			return false;
+		}
+	}
+
+	private void updateIterator() {
+		if(iterator == null || !iterator.hasNext()) {
+			ObjectListing listing = client.listObjects(req);
+			req = req.withMarker(listing.getNextMarker());
+			empty = !listing.isTruncated();
+			iterator =  listing.getObjectSummaries().iterator();
+		}
+	}
+		
+	@Override
+	public S3ObjectSummary next() {
+		updateIterator();
+		return iterator.next();
+	}
+
+}
diff --git a/micro-s3/src/main/java/com/aol/micro/server/s3/S3Utils.java b/micro-s3/src/main/java/com/aol/micro/server/s3/S3Utils.java
@@ -1,29 +1,50 @@
 package com.aol.micro.server.s3;
 
+import java.io.ByteArrayInputStream;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.FileSystems;
+import java.nio.file.Files;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
+import java.util.function.Supplier;
 import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
+import org.apache.commons.io.FileUtils;
+import org.jooq.lambda.tuple.Tuple;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 
+import com.amazonaws.AmazonClientException;
+import com.amazonaws.AmazonServiceException;
 import com.amazonaws.services.s3.AmazonS3Client;
 import com.amazonaws.services.s3.model.DeleteObjectsRequest;
 import com.amazonaws.services.s3.model.DeleteObjectsRequest.KeyVersion;
 import com.amazonaws.services.s3.model.ListObjectsRequest;
 import com.amazonaws.services.s3.model.ObjectListing;
 import com.amazonaws.services.s3.model.S3ObjectSummary;
+import com.amazonaws.services.s3.transfer.Download;
+import com.amazonaws.services.s3.transfer.TransferManager;
 import com.aol.cyclops.control.ReactiveSeq;
+import com.aol.cyclops.util.ExceptionSoftener;
 
 @Component
 public class S3Utils {
 
 	private final AmazonS3Client client;
+	private final TransferManager transferManager;
+	private final String tmpDirectory;
 
 	@Autowired
-	public S3Utils(AmazonS3Client client) {
+	public S3Utils(AmazonS3Client client, TransferManager transferManager,
+			@Value("${s3.tmp.dir:#{null}}") String tmpDirectory) {
 		this.client = client;
+		this.transferManager = transferManager;
+		this.tmpDirectory = tmpDirectory;
 	}
 
 	public List<S3ObjectSummary> getAllSummaries(ListObjectsRequest req) {
@@ -36,14 +57,13 @@ public List<S3ObjectSummary> getAllSummaries(ListObjectsRequest req) {
 			marker = listing.getNextMarker();
 			result.addAll(listing.getObjectSummaries());
 		} while (listing.isTruncated());
+
 		return result;
 	}
 
-	/*
-	 * TODO implement smarter mechanism to reduce number of queries
-	 */
 	public <T> Stream<T> getSummariesStream(ListObjectsRequest req, Function<S3ObjectSummary, T> processor) {
-		return getAllSummaries(req).stream().map(processor);
+		Iterable<S3ObjectSummary> iterable = () -> new S3ObjectSummaryIterator(client, req);
+		return StreamSupport.stream(iterable.spliterator(), false).map(processor);
 	}
 
 	public void delete(String bucketName, List<KeyVersion> objects) {
@@ -54,4 +74,23 @@ public void delete(String bucketName, List<KeyVersion> objects) {
 		});
 	}
 
+	public InputStream getInputStream(String bucketName, String key, Supplier<File> tempFileSupplier)
+			throws AmazonServiceException, AmazonClientException, InterruptedException, IOException {
+		File file = tempFileSupplier.get();
+		try {
+			Download download = transferManager.download(bucketName, key, file);
+			download.waitForCompletion();
+			return new ByteArrayInputStream(FileUtils.readFileToByteArray(file));
+		} finally {
+			file.delete();
+		}
+	}
+
+	public InputStream getInputStream(String bucketName, String key)
+			throws AmazonServiceException, AmazonClientException, InterruptedException, IOException {
+		Supplier<File> tempFileSupplier = ExceptionSoftener.softenSupplier(() -> Files
+				.createTempFile(FileSystems.getDefault().getPath(tmpDirectory), "micro-s3", "file").toFile());
+		return getInputStream(bucketName, key, tempFileSupplier);
+	}
+
 }
