Add config to set the BigQuery Job timeout (GoogleCloudDataproc#1163)

README-template.md

@@ -885,6 +885,15 @@ word-break:break-word
      </td>
      <td>Read</td>
    </tr>
+  <tr>
+     <td><code>bigQueryJobTimeoutInMinutes</code>
+     </td>
+     <td>Config to set the BigQuery job timeout in minutes.
+          Default value is <code>360</code> minutes.
+         <br/> (Optional)
+     </td>
+     <td>Read/Write</td>
+   </tr>
 
 </table>
 

bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryClient.java

@@ -85,6 +85,7 @@ public class BigQueryClient {
   private final Optional<String> materializationDataset;
   private final JobConfigurationFactory jobConfigurationFactory;
   private final Optional<BigQueryJobCompletionListener> jobCompletionListener;
+  private final long bigQueryJobTimeoutInMinutes;
 
   public BigQueryClient(
       BigQuery bigQuery,
@@ -93,13 +94,15 @@ public BigQueryClient(
       Cache<String, TableInfo> destinationTableCache,
       Map<String, String> labels,
       Priority queryJobPriority,
-      Optional<BigQueryJobCompletionListener> jobCompletionListener) {
+      Optional<BigQueryJobCompletionListener> jobCompletionListener,
+      long bigQueryJobTimeoutInMinutes) {
     this.bigQuery = bigQuery;
     this.materializationProject = materializationProject;
     this.materializationDataset = materializationDataset;
     this.destinationTableCache = destinationTableCache;
     this.jobConfigurationFactory = new JobConfigurationFactory(labels, queryJobPriority);
     this.jobCompletionListener = jobCompletionListener;
+    this.bigQueryJobTimeoutInMinutes = bigQueryJobTimeoutInMinutes;
   }
 
   public static synchronized void runCleanupJobs() {
@@ -127,10 +130,16 @@ public JobInfo waitForJob(Job job) {
       Job completedJob =
           job.waitFor(
               RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
-              RetryOption.totalTimeout(Duration.ofMinutes(3)));
-      if (completedJob == null && completedJob.getStatus().getError() != null) {
+              RetryOption.totalTimeout(Duration.ofMinutes(bigQueryJobTimeoutInMinutes)));
+      if (completedJob == null || completedJob.getStatus().getError() != null) {
         throw new UncheckedIOException(
-            new IOException(completedJob.getStatus().getError().toString()));
+            new IOException(
+                completedJob != null ? completedJob.getStatus().getError().toString() : null));
+      }
+      if (!completedJob.isDone()) {
+        completedJob.cancel();
+        throw new IllegalStateException(
+            String.format("Job aborted due to timeout  : %s minutes", bigQueryJobTimeoutInMinutes));
       }
       jobCompletionListener.ifPresent(jcl -> jcl.accept(completedJob));
       return completedJob;

bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryClientFactoryConfig.java

@@ -53,8 +53,9 @@ public class BigQueryClientFactoryConfig implements BigQueryConfig {
   private final int channelPoolSize;
   private final Optional<Integer> flowControlWindowBytes;
   private final QueryJobConfiguration.Priority queryJobPriority;
+  private final long bigQueryJobTimeoutInMinutes;
 
-  BigQueryClientFactoryConfig(BigQueryConfig bigQueryConfig) {
+  BigQueryClientFactoryConfig(BigQueryConfig bigQueryConfig, long bigQueryJobTimeoutInMinutes) {
     this.accessTokenProviderFQCN = bigQueryConfig.getAccessTokenProviderFQCN();
     this.accessTokenProviderConfig = bigQueryConfig.getAccessTokenProviderConfig();
     this.credentialsKey = bigQueryConfig.getCredentialsKey();
@@ -85,6 +86,7 @@ public class BigQueryClientFactoryConfig implements BigQueryConfig {
     this.channelPoolSize = bigQueryConfig.getChannelPoolSize();
     this.flowControlWindowBytes = bigQueryConfig.getFlowControlWindowBytes();
     this.queryJobPriority = bigQueryConfig.getQueryJobPriority();
+    this.bigQueryJobTimeoutInMinutes = bigQueryJobTimeoutInMinutes;
   }
 
   @Override
@@ -222,6 +224,10 @@ public Priority getQueryJobPriority() {
     return queryJobPriority;
   }
 
+  public long getBigQueryJobTimeoutInMinutes() {
+    return bigQueryJobTimeoutInMinutes;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryClientModule.java

@@ -137,6 +137,7 @@ public BigQueryClient provideBigQueryClient(
         destinationTableCache,
         bigQueryJobLabels,
         config.getQueryJobPriority(),
-        Optional.of(jobCompletionListener));
+        Optional.of(jobCompletionListener),
+        config.getBigQueryJobTimeoutInMinutes());
   }
 }

bigquery-connector-common/src/main/java/com/google/cloud/bigquery/connector/common/BigQueryConfig.java

@@ -81,6 +81,8 @@ public interface BigQueryConfig {
 
   Priority getQueryJobPriority();
 
+  long getBigQueryJobTimeoutInMinutes();
+
   default int getClientCreationHashCode() {
     return Objects.hashCode(
         getAccessTokenProviderFQCN(),

bigquery-connector-common/src/test/java/com/google/cloud/bigquery/connector/common/BigQueryClientFactoryTest.java

@@ -600,6 +600,11 @@ public Priority getQueryJobPriority() {
       return Priority.INTERACTIVE;
     }
 
+    @Override
+    public long getBigQueryJobTimeoutInMinutes() {
+      return 6 * 60;
+    }
+
     @Override
     public boolean equals(Object o) {
       if (this == o) {

spark-bigquery-connector-common/src/main/java/com/google/cloud/spark/bigquery/SparkBigQueryConfig.java

@@ -152,6 +152,8 @@ public static WriteMethod from(@Nullable String writeMethod) {
       "spark.sql.sources.partitionOverwriteMode";
 
   public PartitionOverwriteMode partitionOverwriteModeValue = PartitionOverwriteMode.STATIC;
+  public static final String BIGQUERY_JOB_TIMEOUT_IN_MINUTES = "bigQueryJobTimeoutInMinutes";
+  static final long BIGQUERY_JOB_TIMEOUT_IN_MINUTES_DEFAULT = 6 * 60; // 6 hrs
 
   TableId tableId;
   // as the config needs to be Serializable, internally it uses
@@ -227,6 +229,7 @@ public static WriteMethod from(@Nullable String writeMethod) {
   private com.google.common.base.Optional<String> destinationTableKmsKeyName = empty();
 
   private boolean allowMapTypeConversion = ALLOW_MAP_TYPE_CONVERSION_DEFAULT;
+  private long bigQueryJobTimeoutInMinutes = BIGQUERY_JOB_TIMEOUT_IN_MINUTES_DEFAULT;
 
   @VisibleForTesting
   SparkBigQueryConfig() {
@@ -569,6 +572,11 @@ public static SparkBigQueryConfig from(
             .transform(PartitionOverwriteMode::valueOf)
             .or(PartitionOverwriteMode.STATIC);
 
+    config.bigQueryJobTimeoutInMinutes =
+        getAnyOption(globalOptions, options, BIGQUERY_JOB_TIMEOUT_IN_MINUTES)
+            .transform(Long::valueOf)
+            .or(BIGQUERY_JOB_TIMEOUT_IN_MINUTES_DEFAULT);
+
     return config;
   }
 
@@ -1009,6 +1017,10 @@ public boolean getAllowMapTypeConversion() {
     return allowMapTypeConversion;
   }
 
+  public long getBigQueryJobTimeoutInMinutes() {
+    return bigQueryJobTimeoutInMinutes;
+  }
+
   public ImmutableMap<String, String> getBigQueryTableLabels() {
     return bigQueryTableLabels;
   }

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/SparkBigQueryConfigTest.java

@@ -130,6 +130,7 @@ public void testDefaults() {
     assertThat(config.getQueryJobPriority()).isEqualTo(SparkBigQueryConfig.DEFAULT_JOB_PRIORITY);
     assertThat(config.getKmsKeyName()).isEqualTo(Optional.empty());
     assertThat(config.getAllowMapTypeConversion()).isTrue();
+    assertThat(config.getBigQueryJobTimeoutInMinutes()).isEqualTo(6 * 60);
   }
 
   @Test
@@ -176,6 +177,7 @@ public void testConfigFromOptions() {
                 .put("queryJobPriority", "batch")
                 .put("destinationTableKmsKeyName", "some/key/name")
                 .put("allowMapTypeConversion", "false")
+                .put("bigQueryJobTimeoutInMinutes", "30")
                 .build());
     SparkBigQueryConfig config =
         SparkBigQueryConfig.from(
@@ -229,6 +231,7 @@ public void testConfigFromOptions() {
     assertThat(config.getQueryJobPriority()).isEqualTo(Priority.valueOf("BATCH"));
     assertThat(config.getKmsKeyName()).isEqualTo(Optional.of("some/key/name"));
     assertThat(config.getAllowMapTypeConversion()).isFalse();
+    assertThat(config.getBigQueryJobTimeoutInMinutes()).isEqualTo(30);
   }
 
   @Test

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/integration/IntegrationTestUtils.java

@@ -66,7 +66,8 @@ private static BigQueryClient getBigQueryClient() {
         destinationTableCache,
         ImmutableMap.of(),
         SparkBigQueryConfig.DEFAULT_JOB_PRIORITY,
-        Optional.empty());
+        Optional.empty(),
+        6 * 60);
   }
 
   public static void createDataset(String dataset) {

spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/integration/ReadFromQueryIntegrationTestBase.java

@@ -208,6 +208,27 @@ public void testQueryJobPriority() {
 
     validateResult(df);
   }
+
+  @Test
+  public void testReadFromLongQueryWithBigQueryJobTimeout() {
+    String query = "SELECT * FROM `largesamples.wikipedia_pageviews_201001`";
+    assertThrows(
+        RuntimeException.class,
+        () -> {
+          try {
+            spark
+                .read()
+                .format("bigquery")
+                .option("viewsEnabled", true)
+                .option("materializationDataset", testDataset.toString())
+                .option("bigQueryJobTimeoutInMinutes", "1")
+                .load(query)
+                .show();
+          } catch (Exception e) {
+            throw e;
+          }
+        });
+  }
 }
 
 class TestBigQueryJobCompletionListener extends SparkListener {