From 4be06be82d053311b91dab74526581fd574e047a Mon Sep 17 00:00:00 2001
From: Andrew Sung <asung@perfectsensedigital.com>
Date: Fri, 10 Jul 2020 16:46:15 -0400
Subject: [PATCH] Adds storage objects iterator.

---
 .../gyro/google/GoogleStorageFileBackend.java | 35 ++-------
 .../gyro/google/StorageObjectIterator.java    | 74 +++++++++++++++++++
 2 files changed, 82 insertions(+), 27 deletions(-)
 create mode 100644 src/main/java/gyro/google/StorageObjectIterator.java

diff --git a/src/main/java/gyro/google/GoogleStorageFileBackend.java b/src/main/java/gyro/google/GoogleStorageFileBackend.java
index 5171d9e7..f6037386 100644
--- a/src/main/java/gyro/google/GoogleStorageFileBackend.java
+++ b/src/main/java/gyro/google/GoogleStorageFileBackend.java
@@ -19,13 +19,12 @@
 import java.io.ByteArrayOutputStream;
 import java.io.InputStream;
 import java.io.OutputStream;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.Spliterators;
 import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 import com.google.api.client.http.ByteArrayContent;
 import com.google.api.services.storage.Storage;
-import com.google.api.services.storage.model.Objects;
 import com.google.api.services.storage.model.StorageObject;
 import com.psddev.dari.util.ObjectUtils;
 import gyro.core.FileBackend;
@@ -72,27 +71,9 @@ public String getCredentials() {
     @Override
     public Stream<String> list() throws Exception {
         if (this.equals(GyroCore.getStateBackend(getName()))) {
-            List<StorageObject> storageObjects = new ArrayList<>();
-            Storage client = client();
-            String nextPageToken = null;
-            int count = 0;
-
-            do {
-                Objects objects = client.objects()
-                    .list(getBucket())
-                    .setPrefix(prefixed(""))
-                    .setPageToken(nextPageToken)
-                    .execute();
-
-                if (objects.getItems() != null) {
-                    storageObjects.addAll(objects.getItems());
-                }
-
-                nextPageToken = objects.getNextPageToken();
-                count++;
-            } while (nextPageToken != null || count < 10);
-
-            return storageObjects.stream()
+            return StreamSupport.stream(
+                Spliterators.spliteratorUnknownSize(new StorageObjectIterator(getBucket(), prefixed(""), client()), 0),
+                false)
                 .map(StorageObject::getName)
                 .filter(f -> f.endsWith(".gyro"))
                 .map(this::removePrefix);
@@ -109,6 +90,7 @@ public InputStream openInput(String file) throws Exception {
     @Override
     public OutputStream openOutput(String file) throws Exception {
         return new ByteArrayOutputStream() {
+
             public void close() {
                 try {
                     StorageObject upload = new StorageObject();
@@ -128,8 +110,8 @@ public void delete(String file) throws Exception {
 
     private Storage client() {
         GoogleCredentials credentials = (GoogleCredentials) getRootScope().getSettings(CredentialsSettings.class)
-                .getCredentialsByName()
-                .get("google::" + getCredentials());
+            .getCredentialsByName()
+            .get("google::" + getCredentials());
 
         return credentials.createClient(Storage.class);
     }
@@ -145,5 +127,4 @@ private String removePrefix(String file) {
 
         return file;
     }
-
 }
diff --git a/src/main/java/gyro/google/StorageObjectIterator.java b/src/main/java/gyro/google/StorageObjectIterator.java
new file mode 100644
index 00000000..adf6ecd7
--- /dev/null
+++ b/src/main/java/gyro/google/StorageObjectIterator.java
@@ -0,0 +1,74 @@
+package gyro.google;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+import com.google.api.services.storage.Storage;
+import com.google.api.services.storage.model.Objects;
+import com.google.api.services.storage.model.StorageObject;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class StorageObjectIterator implements Iterator<StorageObject> {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(StorageObjectIterator.class);
+
+    private String bucket;
+    private String prefix;
+    private Storage client;
+
+    private List<StorageObject> storageObjects = new ArrayList<>();
+    private int index;
+    private String nextPageToken;
+
+    public StorageObjectIterator(String bucket, String prefix, Storage client) {
+        this.bucket = bucket;
+        this.prefix = prefix;
+        this.client = client;
+    }
+
+    @Override
+    public boolean hasNext() {
+        if (index < storageObjects.size()) {
+            return true;
+        }
+
+        if (index > 0 && nextPageToken == null) {
+            return false;
+        }
+
+        try {
+            Objects objects = client.objects()
+                .list(bucket)
+                .setPrefix(prefix)
+                .setPageToken(nextPageToken)
+                .setMaxResults(100L)
+                .execute();
+
+            if (objects.getItems() != null) {
+                index = 0;
+                storageObjects.clear();
+                storageObjects.addAll(objects.getItems());
+            }
+            nextPageToken = objects.getNextPageToken();
+
+        } catch (IOException e) {
+            LOGGER.error("Failed to retrieve storage objects: {} {} {}", bucket, prefix, nextPageToken, e);
+            return false;
+        }
+
+        return !storageObjects.isEmpty();
+
+    }
+
+    @Override
+    public StorageObject next() {
+        if (hasNext()) {
+            return storageObjects.get(index++);
+        }
+        throw new NoSuchElementException();
+    }
+}