diff --git a/docs/src/main/sphinx/connector/googlesheets.rst b/docs/src/main/sphinx/connector/googlesheets.rst index 8ef390544f4b..25475821a46c 100644 --- a/docs/src/main/sphinx/connector/googlesheets.rst +++ b/docs/src/main/sphinx/connector/googlesheets.rst @@ -33,6 +33,7 @@ Property name Description ``gsheets.metadata-sheet-id`` Sheet ID of the spreadsheet, that contains the table mapping ``gsheets.max-data-cache-size`` Maximum number of spreadsheets to cache, defaults to ``1000`` ``gsheets.data-cache-ttl`` How long to cache spreadsheet data or metadata, defaults to ``5m`` +``gsheets.read-timeout`` Timeout to read data from spreadsheet, defaults to ``20s`` =================================== ===================================================================== Credentials 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 e6888a81e540..87299fb27880 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,6 +15,7 @@ import com.google.api.client.auth.oauth2.Credential; import com.google.api.client.googleapis.auth.oauth2.GoogleCredential; +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.services.sheets.v4.Sheets; @@ -27,6 +28,7 @@ import com.google.common.util.concurrent.UncheckedExecutionException; import io.airlift.json.JsonCodec; import io.airlift.log.Logger; +import io.airlift.units.Duration; import io.trino.collect.cache.NonEvictableLoadingCache; import io.trino.spi.TrinoException; import io.trino.spi.type.VarcharType; @@ -52,6 +54,7 @@ import static io.trino.plugin.google.sheets.SheetsErrorCode.SHEETS_METASTORE_ERROR; import static io.trino.plugin.google.sheets.SheetsErrorCode.SHEETS_TABLE_LOAD_ERROR; import static io.trino.plugin.google.sheets.SheetsErrorCode.SHEETS_UNKNOWN_TABLE_ERROR; +import static java.lang.Math.toIntExact; import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; import static java.util.concurrent.TimeUnit.MILLISECONDS; @@ -82,7 +85,7 @@ public SheetsClient(SheetsConfig config, JsonCodec this.credentialsFilePath = config.getCredentialsFilePath(); try { - this.sheetsService = new Sheets.Builder(newTrustedTransport(), JSON_FACTORY, getCredentials()).setApplicationName(APPLICATION_NAME).build(); + this.sheetsService = new Sheets.Builder(newTrustedTransport(), JSON_FACTORY, setTimeout(getCredentials(), config.getReadTimeout())).setApplicationName(APPLICATION_NAME).build(); } catch (GeneralSecurityException | IOException e) { throw new TrinoException(SHEETS_BAD_CREDENTIALS_ERROR, e); @@ -230,4 +233,13 @@ private static CacheBuilder newCacheBuilder(long expiresAfterWri { return CacheBuilder.newBuilder().expireAfterWrite(expiresAfterWriteMillis, MILLISECONDS).maximumSize(maximumSize); } + + private static HttpRequestInitializer setTimeout(HttpRequestInitializer requestInitializer, Duration readTimeout) + { + requireNonNull(readTimeout, "readTimeout is null"); + return httpRequest -> { + requestInitializer.initialize(httpRequest); + httpRequest.setReadTimeout(toIntExact(readTimeout.toMillis())); + }; + } } diff --git a/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsConfig.java b/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsConfig.java index 48fa3a33039f..8299d9ab3f09 100644 --- a/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsConfig.java +++ b/plugin/trino-google-sheets/src/main/java/io/trino/plugin/google/sheets/SheetsConfig.java @@ -31,6 +31,7 @@ public class SheetsConfig private String metadataSheetId; private int sheetsDataMaxCacheSize = 1000; private Duration sheetsDataExpireAfterWrite = new Duration(5, TimeUnit.MINUTES); + private Duration readTimeout = new Duration(20, TimeUnit.SECONDS); // 20s is the default timeout of com.google.api.client.http.HttpRequest @NotNull @FileExists @@ -92,4 +93,17 @@ public SheetsConfig setSheetsDataExpireAfterWrite(Duration sheetsDataExpireAfter this.sheetsDataExpireAfterWrite = sheetsDataExpireAfterWriteMinutes; return this; } + + @MinDuration("0ms") + public Duration getReadTimeout() + { + return readTimeout; + } + + @Config("gsheets.read-timeout") + public SheetsConfig setReadTimeout(Duration readTimeout) + { + this.readTimeout = readTimeout; + return this; + } } 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 0b2f2046ae17..2eb742e46d29 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 @@ -28,7 +28,7 @@ public class TestGoogleSheets protected QueryRunner createQueryRunner() throws Exception { - return createSheetsQueryRunner(ImmutableMap.of(), ImmutableMap.of()); + return createSheetsQueryRunner(ImmutableMap.of(), ImmutableMap.of("gsheets.read-timeout", "1m")); } @Test diff --git a/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestSheetsConfig.java b/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestSheetsConfig.java index 5a4e0441cdd6..9988a050abf0 100644 --- a/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestSheetsConfig.java +++ b/plugin/trino-google-sheets/src/test/java/io/trino/plugin/google/sheets/TestSheetsConfig.java @@ -37,7 +37,8 @@ public void testDefaults() .setCredentialsFilePath(null) .setMetadataSheetId(null) .setSheetsDataMaxCacheSize(1000) - .setSheetsDataExpireAfterWrite(new Duration(5, TimeUnit.MINUTES))); + .setSheetsDataExpireAfterWrite(new Duration(5, TimeUnit.MINUTES)) + .setReadTimeout(new Duration(20, TimeUnit.SECONDS))); } @Test @@ -51,13 +52,15 @@ public void testExplicitPropertyMappings() .put("gsheets.metadata-sheet-id", "foo_bar_sheet_id#Sheet1") .put("gsheets.max-data-cache-size", "2000") .put("gsheets.data-cache-ttl", "10m") + .put("gsheets.read-timeout", "1m") .buildOrThrow(); SheetsConfig expected = new SheetsConfig() .setCredentialsFilePath(credentialsFile.toString()) .setMetadataSheetId("foo_bar_sheet_id#Sheet1") .setSheetsDataMaxCacheSize(2000) - .setSheetsDataExpireAfterWrite(new Duration(10, TimeUnit.MINUTES)); + .setSheetsDataExpireAfterWrite(new Duration(10, TimeUnit.MINUTES)) + .setReadTimeout(new Duration(1, TimeUnit.MINUTES)); assertFullMapping(properties, expected); }