From bac4e22bff2e41490a8dcfb89100dcdb6dbde03d Mon Sep 17 00:00:00 2001 From: Juan Celhay Date: Tue, 3 Sep 2024 10:27:00 -0400 Subject: [PATCH] Add retries to DriveConnection.listFiles() on GoogleJsonResponseExceptions (#2541) * Add retries to DriveConnection.listFiles() on GoogleJsonResponseExceptions * Remove unused import * Remove unread variable * Add comment inputs * fix formatting * Remove period from error message. --- .../storage/drive/DriveConnection.java | 36 ++++++++++++-- .../storage/drive/DriveConnectionTest.java | 48 +++++++++++++++++++ 2 files changed, 79 insertions(+), 5 deletions(-) diff --git a/core/src/main/java/google/registry/storage/drive/DriveConnection.java b/core/src/main/java/google/registry/storage/drive/DriveConnection.java index e4fbc7c86..c7f914963 100644 --- a/core/src/main/java/google/registry/storage/drive/DriveConnection.java +++ b/core/src/main/java/google/registry/storage/drive/DriveConnection.java @@ -14,6 +14,7 @@ package google.registry.storage.drive; +import com.google.api.client.googleapis.json.GoogleJsonResponseException; import com.google.api.client.http.ByteArrayContent; import com.google.api.services.drive.Drive; import com.google.api.services.drive.Drive.Files; @@ -21,6 +22,7 @@ import com.google.api.services.drive.model.File; import com.google.api.services.drive.model.FileList; import com.google.common.base.Strings; import com.google.common.collect.ImmutableList; +import com.google.common.flogger.FluentLogger; import com.google.common.net.MediaType; import java.io.IOException; import java.util.List; @@ -30,6 +32,11 @@ import javax.inject.Inject; /** Class encapsulating parameters and state for accessing the Drive API. */ public class DriveConnection { + private static final FluentLogger logger = FluentLogger.forEnclosingClass(); + + /** Number of times a request to Drive will be retried before propagating a failure. */ + private static final int MAX_RETRIES = 10; + private static final MediaType GOOGLE_FOLDER = MediaType.create("application", "vnd.google-apps.folder"); @@ -76,9 +83,9 @@ public class DriveConnection { * Creates a file with the given parent or updates the existing one if a file already exists with * that same name and parent. * + * @return the file id. * @throws IllegalStateException if multiple files with that name exist in the given folder. * @throws IOException if communication with Google Drive fails for any reason. - * @return the file id. */ public String createOrUpdateFile( String name, MediaType mimeType, String parentFolderId, byte[] bytes) throws IOException { @@ -126,6 +133,8 @@ public class DriveConnection { * @see The query format */ public List listFiles(String parentFolderId, String query) throws IOException { + int failures = 0; + boolean latestRequestFailed; ImmutableList.Builder result = new ImmutableList.Builder<>(); Files.List req = drive.files().list(); StringBuilder q = new StringBuilder(String.format("'%s' in parents", parentFolderId)); @@ -134,10 +143,27 @@ public class DriveConnection { } req.setQ(q.toString()); do { - FileList files = req.execute(); - files.getFiles().forEach(file -> result.add(file.getId())); - req.setPageToken(files.getNextPageToken()); - } while (!Strings.isNullOrEmpty(req.getPageToken())); + try { + latestRequestFailed = false; + FileList files = req.execute(); + files.getFiles().forEach(file -> result.add(file.getId())); + req.setPageToken(files.getNextPageToken()); + } catch (GoogleJsonResponseException e) { + if (failures >= MAX_RETRIES) { + throw new RuntimeException( + String.format( + "Max failures reached while attempting to list Drive files in folder %s with " + + "query %s; failing permanently.", + parentFolderId, query), + e); + } + latestRequestFailed = true; + logger.atWarning().withCause(e).log( + "Attempt: %d. Failed to list files from Drive. Folder: %s, query: %s.", + failures, parentFolderId, query); + failures++; + } + } while (!Strings.isNullOrEmpty(req.getPageToken()) || latestRequestFailed); return result.build(); } diff --git a/core/src/test/java/google/registry/storage/drive/DriveConnectionTest.java b/core/src/test/java/google/registry/storage/drive/DriveConnectionTest.java index a5003cb27..a076472ed 100644 --- a/core/src/test/java/google/registry/storage/drive/DriveConnectionTest.java +++ b/core/src/test/java/google/registry/storage/drive/DriveConnectionTest.java @@ -25,7 +25,11 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import com.google.api.client.googleapis.json.GoogleJsonError; +import com.google.api.client.googleapis.json.GoogleJsonResponseException; import com.google.api.client.http.ByteArrayContent; +import com.google.api.client.http.HttpHeaders; +import com.google.api.client.http.HttpResponseException; import com.google.api.services.drive.Drive; import com.google.api.services.drive.Drive.Files; import com.google.api.services.drive.model.File; @@ -148,6 +152,50 @@ class DriveConnectionTest { verify(filesList, times(2)).getPageToken(); } + @Test + void testListFiles_succeedsRetriedGoogleJsonResponseException() throws Exception { + GoogleJsonResponseException googleJsonResponseException = + new GoogleJsonResponseException( + new HttpResponseException.Builder(503, "Service Unavailable.", new HttpHeaders()), + new GoogleJsonError()); + File testFile = new File().setId("testFile"); + File testFile1 = new File().setId("testFile1"); + + when(filesList.execute()) + .thenThrow(googleJsonResponseException) + .thenReturn(new FileList().setFiles(ImmutableList.of(testFile)).setNextPageToken("next")) + .thenThrow(googleJsonResponseException) + .thenReturn(new FileList().setFiles(ImmutableList.of(testFile1))); + + assertThat(driveConnection.listFiles("driveFolderId")) + .containsExactly(testFile.getId(), testFile1.getId()); + verify(filesList).setPageToken("next"); + verify(filesList, times(1)).setQ("'driveFolderId' in parents"); + verify(filesList, times(4)).getPageToken(); + } + + @Test + void testListFiles_throwsException_afterMaxRetriedGoogleJsonResponseException() throws Exception { + GoogleJsonResponseException googleJsonResponseException = + new GoogleJsonResponseException( + new HttpResponseException.Builder(503, "Service Unavailable.", new HttpHeaders()), + new GoogleJsonError()); + when(filesList.execute()).thenThrow(googleJsonResponseException); + + Exception thrown = + assertThrows( + Exception.class, () -> driveConnection.listFiles("driveFolderId", "sampleQuery")); + assertThat(thrown.getCause()).isEqualTo(googleJsonResponseException); + assertThat(thrown.getMessage()) + .isEqualTo( + "Max failures reached while attempting to list Drive files in folder " + + "driveFolderId with query sampleQuery; failing permanently."); + + verify(filesList, times(0)).setPageToken(null); + verify(filesList, times(1)).setQ("'driveFolderId' in parents and sampleQuery"); + verify(filesList, times(10)).getPageToken(); + } + @Test void testCreateOrUpdateFile_succeedsForNewFile() throws Exception { when(files.create(