Adding read timeout to BQ Execute Config with default 120s

Author: masivesky

src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryArgumentSetter.java

@@ -80,7 +80,7 @@ public void run(ActionContext context) throws Exception {
     Credentials credentials = config.getServiceAccount() == null ?
       null : GCPUtils.loadServiceAccountCredentials(config.getServiceAccount(),
                                                     config.isServiceAccountFilePath());
-    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials);
+    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials, null);
     Job queryJob = bigQuery.create(JobInfo.newBuilder(queryConfig).setJobId(jobId).build());
 
     LOG.info("Executing SQL as job {}.", jobId.getJob());

src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecute.java

@@ -125,7 +125,7 @@ public void run(ActionContext context) throws Exception {
     Credentials credentials = config.getServiceAccount() == null ?
       null : GCPUtils.loadServiceAccountCredentials(config.getServiceAccount(),
                                                     config.isServiceAccountFilePath());
-    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials);
+    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials, config.getReadTimeout());
     //create dataset to store the results if not exists
     if (config.getStoreResults() && !Strings.isNullOrEmpty(datasetName) &&
       !Strings.isNullOrEmpty(tableName)) {
@@ -283,11 +283,13 @@ public static final class Config extends AbstractBigQueryActionConfig {
     private static final String NAME_MAX_RETRY_DURATION = "maxRetryDuration";
     private static final String NAME_RETRY_MULTIPLIER = "retryMultiplier";
     private static final String NAME_MAX_RETRY_COUNT = "maxRetryCount";
+    private static final String NAME_READ_TIMEOUT = "readTimeout";
     public static final long DEFAULT_INITIAL_RETRY_DURATION_SECONDS = 1L;
     public static final double DEFAULT_RETRY_MULTIPLIER = 2.0;
     public static final int DEFAULT_MAX_RETRY_COUNT = 5;
     // Sn = a * (1 - r^n) / (r - 1)
     public static final long DEFULT_MAX_RETRY_DURATION_SECONDS = 63L;
+    public static final int DEFAULT_READ_TIMEOUT = 120;
 
     @Description("Dialect of the SQL command. The value must be 'legacy' or 'standard'. " +
       "If set to 'standard', the query will use BigQuery's standard SQL: " +
@@ -390,13 +392,19 @@ public static final class Config extends AbstractBigQueryActionConfig {
             "are reserved keys and cannot be used as label keys.")
     protected String jobLabelKeyValue;
 
+    @Name(NAME_READ_TIMEOUT)
+    @Nullable
+    @Macro
+    @Description("Timeout in seconds to read data from an established HTTP connection (Default value is 120).")
+    private Integer readTimeout;
+
     private Config(@Nullable String project, @Nullable String serviceAccountType, @Nullable String serviceFilePath,
                    @Nullable String serviceAccountJson, @Nullable String dataset, @Nullable String table,
                    @Nullable String location, @Nullable String cmekKey, @Nullable String dialect, @Nullable String sql,
                    @Nullable String mode, @Nullable Boolean storeResults, @Nullable String jobLabelKeyValue,
                    @Nullable String rowAsArguments, @Nullable Boolean retryOnBackendError,
                    @Nullable Long initialRetryDuration, @Nullable Long maxRetryDuration,
-                   @Nullable Double retryMultiplier, @Nullable Integer maxRetryCount) {
+                   @Nullable Double retryMultiplier, @Nullable Integer maxRetryCount, @Nullable Integer readTimeout) {
       this.project = project;
       this.serviceAccountType = serviceAccountType;
       this.serviceFilePath = serviceFilePath;
@@ -416,6 +424,7 @@ private Config(@Nullable String project, @Nullable String serviceAccountType, @N
       this.maxRetryDuration = maxRetryDuration;
       this.maxRetryCount = maxRetryCount;
       this.retryMultiplier = retryMultiplier;
+      this.readTimeout = readTimeout;
     }
 
     public boolean isLegacySQL() {
@@ -481,6 +490,10 @@ public int getMaxRetryCount() {
       return maxRetryCount == null ? DEFAULT_MAX_RETRY_COUNT : maxRetryCount;
     }
 
+    public int getReadTimeout() {
+      return readTimeout == null ? DEFAULT_READ_TIMEOUT : readTimeout;
+    }
+
     @Override
     public void validate(FailureCollector failureCollector) {
       validate(failureCollector, Collections.emptyMap());
@@ -544,14 +557,15 @@ void validateJobLabelKeyValue(FailureCollector failureCollector) {
           !containsMacro(NAME_INITIAL_RETRY_DURATION) && !containsMacro(NAME_MAX_RETRY_DURATION) &&
           !containsMacro(NAME_MAX_RETRY_COUNT) && !containsMacro(NAME_RETRY_MULTIPLIER)) {
           validateRetryConfiguration(
-            failureCollector, initialRetryDuration, maxRetryDuration, maxRetryCount, retryMultiplier
+            failureCollector, initialRetryDuration, maxRetryDuration, maxRetryCount, retryMultiplier, readTimeout
           );
       }
       failureCollector.getOrThrowException();
     }
 
     void validateRetryConfiguration(FailureCollector failureCollector, Long initialRetryDuration,
-                                    Long maxRetryDuration, Integer maxRetryCount, Double retryMultiplier) {
+                                    Long maxRetryDuration, Integer maxRetryCount, Double retryMultiplier,
+                                    Integer readTimeout) {
       if (initialRetryDuration != null && initialRetryDuration <= 0) {
         failureCollector.addFailure("Initial retry duration must be greater than 0.",
                         "Please specify a valid initial retry duration.")
@@ -577,6 +591,11 @@ void validateRetryConfiguration(FailureCollector failureCollector, Long initialR
                         "Please specify a valid max retry duration.")
                 .withConfigProperty(NAME_MAX_RETRY_DURATION);
       }
+      if (readTimeout != null && readTimeout <= 0) {
+        failureCollector.addFailure("Read timeout must be greater than 0.",
+            "Please specify a valid read timeout")
+            .withConfigProperty(NAME_READ_TIMEOUT);
+      }
     }
 
     void validateCmekKey(FailureCollector failureCollector, Map<String, String> arguments) {
@@ -612,7 +631,8 @@ public void validateSQLSyntax(FailureCollector failureCollector, BigQuery bigQue
           } else {
                errorMessage = e.getMessage();
           }
-          failureCollector.addFailure(String.format("%s.", errorMessage), "Please specify a valid query.")
+          failureCollector.addFailure(String.format("%s. Error code: %s.", errorMessage, e.getCode()),
+                        "Please specify a valid query.")
                     .withConfigProperty(SQL);
       }
     }
@@ -626,7 +646,7 @@ private BigQuery getBigQuery(FailureCollector failureCollector) {
         failureCollector.addFailure(e.getMessage(), null);
         failureCollector.getOrThrowException();
       }
-      return GCPUtils.getBigQuery(getProject(), credentials);
+      return GCPUtils.getBigQuery(getProject(), credentials, getReadTimeout());
     }
 
     public static Builder builder() {
@@ -656,6 +676,7 @@ public static class Builder {
       private Long maxRetryDuration;
       private Integer maxRetryCount;
       private Double retryMultiplier;
+      private Integer readTimeout;
 
       public Builder setProject(@Nullable String project) {
         this.project = project;
@@ -752,6 +773,11 @@ public Builder setRetryMultiplier(@Nullable Double retryMultiplier) {
         return this;
       }
 
+      public Builder setReadTimeout(@Nullable Integer readTimeout) {
+        this.readTimeout = readTimeout;
+        return this;
+      }
+
       public Config build() {
         return new Config(
           project,
@@ -772,7 +798,8 @@ public Config build() {
           initialRetryDuration,
           maxRetryDuration,
           retryMultiplier,
-          maxRetryCount
+          maxRetryCount,
+          readTimeout
         );
       }
     }

src/main/java/io/cdap/plugin/gcp/bigquery/common/BigQueryBaseConfig.java

@@ -213,7 +213,7 @@ public void validateCmekKeyLocation(@Nullable CryptoKeyName cmekKeyName, @Nullab
         DatasetId datasetId = DatasetId.of(datasetProjectId, datasetName);
         TableId tableId = tableName == null ? null : TableId.of(datasetProjectId, datasetName, tableName);
         Credentials credentials = connection.getCredentials(failureCollector);
-        BigQuery bigQuery = GCPUtils.getBigQuery(connection.getProject(), credentials);
+        BigQuery bigQuery = GCPUtils.getBigQuery(connection.getProject(), credentials, null);
         Storage storage = GCPUtils.getStorage(connection.getProject(), credentials);
         if (bigQuery == null || storage == null) {
             return;

src/main/java/io/cdap/plugin/gcp/bigquery/connector/BigQueryConnector.java

@@ -132,7 +132,7 @@ public void test(ConnectorContext context) throws ValidationException {
     }
 
     try {
-      BigQuery bigQuery = GCPUtils.getBigQuery(config.getDatasetProject(), credentials);
+      BigQuery bigQuery = GCPUtils.getBigQuery(config.getDatasetProject(), credentials, null);
       bigQuery.listDatasets(BigQuery.DatasetListOption.pageSize(1));
     } catch (Exception e) {
       failureCollector.addFailure(String.format("Could not connect to BigQuery: %s", e.getMessage()),
@@ -236,7 +236,7 @@ private BigQuery getBigQuery(String project) throws IOException {
         GCPUtils.loadServiceAccountCredentials(config.getServiceAccount(), config.isServiceAccountFilePath());
     }
     // Here project decides where the BQ job is run and under which the datasets is listed
-    return GCPUtils.getBigQuery(project, credentials);
+    return GCPUtils.getBigQuery(project, credentials, null);
   }
 
   /**

src/main/java/io/cdap/plugin/gcp/bigquery/sink/AbstractBigQuerySink.java

@@ -88,7 +88,7 @@ public final void prepareRun(BatchSinkContext context) throws Exception {
     Credentials credentials = serviceAccount == null ?
       null : GCPUtils.loadServiceAccountCredentials(serviceAccount, config.isServiceAccountFilePath());
     String project = config.getProject();
-    bigQuery = GCPUtils.getBigQuery(project, credentials);
+    bigQuery = GCPUtils.getBigQuery(project, credentials, null);
     FailureCollector collector = context.getFailureCollector();
     CryptoKeyName cmekKeyName = CmekUtils.getCmekKey(config.cmekKey, context.getArguments().asMap(), collector);
     collector.getOrThrowException();

src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQueryOutputFormat.java

@@ -810,6 +810,6 @@ private static BigQuery getBigQuery(Configuration config) throws IOException {
     String projectId = ConfigurationUtil.getMandatoryConfig(config,
         BigQueryConfiguration.PROJECT_ID);
     Credentials credentials = GCPUtils.loadCredentialsFromConf(config);
-    return GCPUtils.getBigQuery(projectId, credentials);
+    return GCPUtils.getBigQuery(projectId, credentials, null);
   }
 }

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySource.java

@@ -136,7 +136,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
     // Create BigQuery client
     String serviceAccount = config.getServiceAccount();
     Credentials credentials = BigQuerySourceUtils.getCredentials(config.getConnection());
-    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials);
+    BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials, null);
     Dataset dataset = bigQuery.getDataset(DatasetId.of(config.getDatasetProject(), config.getDataset()));
     Storage storage = GCPUtils.getStorage(config.getProject(), credentials);
 

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceConfig.java

@@ -186,7 +186,7 @@ void validateCmekKey(FailureCollector collector, Map<String, String> arguments)
     }
     DatasetId datasetId = DatasetId.of(getDatasetProject(), getDataset());
     Credentials credentials = connection.getCredentials(collector);
-    BigQuery bigQuery = GCPUtils.getBigQuery(getProject(), credentials);
+    BigQuery bigQuery = GCPUtils.getBigQuery(getProject(), credentials, null);
     if (bigQuery == null) {
       return;
     }

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySourceUtils.java

@@ -171,7 +171,7 @@ public static void deleteBigQueryTemporaryTable(Configuration configuration, Big
     String temporaryTable = configuration.get(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME);
     try {
       Credentials credentials = getCredentials(config.getConnection());
-      BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials);
+      BigQuery bigQuery = GCPUtils.getBigQuery(config.getProject(), credentials, null);
       bigQuery.delete(TableId.of(config.getDatasetProject(), config.getDataset(), temporaryTable));
       LOG.debug("Deleted temporary table '{}'", temporaryTable);
     } catch (IOException e) {

src/main/java/io/cdap/plugin/gcp/bigquery/sqlengine/BigQuerySQLEngine.java

@@ -156,7 +156,7 @@ public void prepareRun(SQLEngineContext context) throws Exception {
     datasetName = sqlEngineConfig.getDataset();
 
     // Initialize BQ and GCS clients.
-    bigQuery = GCPUtils.getBigQuery(project, credentials);
+    bigQuery = GCPUtils.getBigQuery(project, credentials, null);
     storage = GCPUtils.getStorage(project, credentials);
     Dataset dataset = bigQuery.getDataset(DatasetId.of(datasetProject, datasetName));
     bucket = BigQueryUtil.getStagingBucketName(context.getRuntimeArguments(), sqlEngineConfig.getLocation(),