Support setting read timeout in Google Sheets

Also, set the timeout as 1m in TestGoogleSheets to avoid flaky tests.

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

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<Map<String, List<SheetsTable>
         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<Object, Object> 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()));
+        };
+    }
 }

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;
+    }
 }

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

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);
     }