diff --git a/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsClient.java b/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsClient.java index 9f15afae1827..59562670d07f 100644 --- a/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsClient.java +++ b/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsClient.java @@ -15,9 +15,12 @@ import com.google.api.client.auth.oauth2.Credential; import com.google.api.client.googleapis.auth.oauth2.GoogleCredential; +import com.google.api.client.http.HttpBackOffIOExceptionHandler; +import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler; import com.google.api.client.http.HttpRequestInitializer; import com.google.api.client.json.JsonFactory; import com.google.api.client.json.jackson2.JacksonFactory; +import com.google.api.client.util.ExponentialBackOff; import com.google.api.services.sheets.v4.Sheets; import com.google.api.services.sheets.v4.SheetsScopes; import com.google.api.services.sheets.v4.model.ValueRange; @@ -69,6 +72,13 @@ public class SheetsClient public static final String RANGE_SEPARATOR = "#"; private static final Logger log = Logger.get(SheetsClient.class); + public static final ExponentialBackOff BACKOFF = new ExponentialBackOff.Builder() + .setInitialIntervalMillis(500) + .setMaxIntervalMillis(10_000) + .setMaxElapsedTimeMillis(60_000) + .setMultiplier(1.5) + .build(); + private static final String APPLICATION_NAME = "trino google sheets integration"; private static final JsonFactory JSON_FACTORY = JacksonFactory.getDefaultInstance(); @@ -333,6 +343,8 @@ private HttpRequestInitializer setTimeout(HttpRequestInitializer requestInitiali httpRequest.setConnectTimeout(toIntExact(config.getConnectionTimeout().toMillis())); httpRequest.setReadTimeout(toIntExact(config.getReadTimeout().toMillis())); httpRequest.setWriteTimeout(toIntExact(config.getWriteTimeout().toMillis())); + httpRequest.setUnsuccessfulResponseHandler(new HttpBackOffUnsuccessfulResponseHandler(BACKOFF)); + httpRequest.setIOExceptionHandler(new HttpBackOffIOExceptionHandler(BACKOFF)); }; } } diff --git a/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestGoogleSheets.java b/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestGoogleSheets.java index b2dd60df4414..7929aa69399b 100644 --- a/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestGoogleSheets.java +++ b/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestGoogleSheets.java @@ -14,6 +14,8 @@ package io.trino.plugin.google.sheets; import com.google.api.client.googleapis.auth.oauth2.GoogleCredential; +import com.google.api.client.http.HttpBackOffIOExceptionHandler; +import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler; import com.google.api.client.http.HttpRequestInitializer; import com.google.api.client.json.jackson2.JacksonFactory; import com.google.api.services.sheets.v4.Sheets; @@ -36,6 +38,7 @@ import java.util.concurrent.TimeUnit; import static com.google.api.client.googleapis.javanet.GoogleNetHttpTransport.newTrustedTransport; +import static io.trino.plugin.google.sheets.SheetsClient.BACKOFF; import static io.trino.plugin.google.sheets.TestSheetsPlugin.DATA_SHEET_ID; import static io.trino.plugin.google.sheets.TestSheetsPlugin.getTestCredentialsPath; import static io.trino.testing.assertions.Assert.assertEventually; @@ -329,6 +332,8 @@ private static HttpRequestInitializer setTimeout(HttpRequestInitializer requestI requestInitializer.initialize(httpRequest); httpRequest.setConnectTimeout(toIntExact(TimeUnit.MINUTES.toMillis(1))); httpRequest.setReadTimeout(toIntExact(TimeUnit.MINUTES.toMillis(1))); + httpRequest.setUnsuccessfulResponseHandler(new HttpBackOffUnsuccessfulResponseHandler(BACKOFF)); + httpRequest.setIOExceptionHandler(new HttpBackOffIOExceptionHandler(BACKOFF)); }; } }