com.google.cloud.RetryOption Java Examples

/** Example usage of {@code waitFor()} with checking period and timeout. */
// [TARGET waitFor(RetryOption...)]
public boolean waitForWithOptions() throws InterruptedException {
  try {
    // [START ]
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(1)));
    if (completedJob == null) {
      // job no longer exists
    } else if (completedJob.getStatus().getError() != null) {
      // job failed, handle error
    } else {
      // job completed successfully
    }
    // [END ]
  } catch (BigQueryException e) {
    if (e.getCause() instanceof PollException) {
      return false;
    }
    throw e;
  }
  return true;
}
 

/** Example of copying the table to a destination table. */
// [TARGET copy(String, String, JobOption...)]
// [VARIABLE "my_dataset"]
// [VARIABLE "my_destination_table"]
public Job copy(String datasetName, String tableName) {
  // [START ]
  Job job = table.copy(datasetName, tableName);
  // Wait for the job to complete.
  try {
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(3)));
    if (completedJob != null && completedJob.getStatus().getError() == null) {
      // Job completed successfully
    } else {
      // Handle error case
    }
  } catch (InterruptedException e) {
    // Handle interrupted wait
  }
  // [END ]
  return job;
}
 

/** Example copying the table to a destination table. */
// [TARGET copy(TableId, JobOption...)]
// [VARIABLE "my_dataset"]
// [VARIABLE "my_destination_table"]
public Job copyTableId(String dataset, String tableName) throws BigQueryException {
  // [START bigquery_copy_table]
  TableId destinationId = TableId.of(dataset, tableName);
  JobOption options = JobOption.fields(JobField.STATUS, JobField.USER_EMAIL);
  Job job = table.copy(destinationId, options);
  // Wait for the job to complete.
  try {
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(3)));
    if (completedJob != null && completedJob.getStatus().getError() == null) {
      // Job completed successfully.
    } else {
      // Handle error case.
    }
  } catch (InterruptedException e) {
    // Handle interrupted wait
  }
  // [END bigquery_copy_table]
  return job;
}
 

/** Example extracting data to single Google Cloud Storage file. */
// [TARGET extract(String, String, JobOption...)]
// [VARIABLE "CSV"]
// [VARIABLE "gs://my_bucket/filename.csv"]
public Job extractSingle(String format, String gcsUrl) {
  // [START bigquery_extract_table]
  Job job = table.extract(format, gcsUrl);
  // Wait for the job to complete
  try {
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(3)));
    if (completedJob != null && completedJob.getStatus().getError() == null) {
      // Job completed successfully
    } else {
      // Handle error case
    }
  } catch (InterruptedException e) {
    // Handle interrupted wait
  }
  // [END bigquery_extract_table]
  return job;
}
 

/** Example loading data from a single Google Cloud Storage file. */
// [TARGET load(FormatOptions, String, JobOption...)]
// [VARIABLE "gs://my_bucket/filename.csv"]
public Job loadSingle(String sourceUri) {
  // [START bigquery_load_table_gcs_csv]
  Job job = table.load(FormatOptions.csv(), sourceUri);
  // Wait for the job to complete
  try {
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(3)));
    if (completedJob != null && completedJob.getStatus().getError() == null) {
      // Job completed successfully
    } else {
      // Handle error case
    }
  } catch (InterruptedException e) {
    // Handle interrupted wait
  }
  // [END bigquery_load_table_gcs_csv]
  return job;
}
 

@Test
public void testSuccessfulLoad() throws Exception {
    when(table.exists()).thenReturn(Boolean.TRUE);
    when(bq.create(ArgumentMatchers.isA(JobInfo.class))).thenReturn(job);
    when(bq.writer(ArgumentMatchers.isA(WriteChannelConfiguration.class))).thenReturn(tableDataWriteChannel);
    when(tableDataWriteChannel.getJob()).thenReturn(job);
    when(job.waitFor(ArgumentMatchers.isA(RetryOption.class))).thenReturn(job);
    when(job.getStatus()).thenReturn(jobStatus);
    when(job.getStatistics()).thenReturn(stats);

    when(stats.getCreationTime()).thenReturn(0L);
    when(stats.getStartTime()).thenReturn(1L);
    when(stats.getEndTime()).thenReturn(2L);

    final TestRunner runner = buildNewRunner(getProcessor());
    addRequiredPropertiesToRunner(runner);
    runner.assertValid();

    runner.enqueue("{ \"data\": \"datavalue\" }");

    runner.run();

    runner.assertAllFlowFilesTransferred(PutBigQueryBatch.REL_SUCCESS);
}
 

@Test
public void testFailedLoad() throws Exception {
    when(table.exists()).thenReturn(Boolean.TRUE);
    when(bq.create(ArgumentMatchers.isA(JobInfo.class))).thenReturn(job);
    when(bq.writer(ArgumentMatchers.isA(WriteChannelConfiguration.class))).thenReturn(tableDataWriteChannel);
    when(tableDataWriteChannel.getJob()).thenReturn(job);
    when(job.waitFor(ArgumentMatchers.isA(RetryOption.class))).thenThrow(BigQueryException.class);
    when(job.getStatus()).thenReturn(jobStatus);
    when(job.getStatistics()).thenReturn(stats);

    when(stats.getCreationTime()).thenReturn(0L);
    when(stats.getStartTime()).thenReturn(1L);
    when(stats.getEndTime()).thenReturn(2L);

    final TestRunner runner = buildNewRunner(getProcessor());
    addRequiredPropertiesToRunner(runner);
    runner.assertValid();

    runner.enqueue("{ \"data\": \"datavalue\" }");

    runner.run();

    runner.assertAllFlowFilesTransferred(PutBigQueryBatch.REL_FAILURE);
}
 

private Job waitForJob(Job queryJob) throws InterruptedException {
  Job completedJob =
      queryJob.waitFor(
          RetryOption.initialRetryDelay(Duration.ofSeconds(initialRetryDelaySecs())),
          RetryOption.totalTimeout(Duration.ofSeconds(totalTimeoutSecs())));
  if (completedJob == null) {
    throw Status.INTERNAL.withDescription("Job no longer exists").asRuntimeException();
  } else if (completedJob.getStatus().getError() != null) {
    throw Status.INTERNAL
        .withDescription("Job failed: " + completedJob.getStatus().getError())
        .asRuntimeException();
  }
  return completedJob;
}
 

@Test
public void shouldReturnTableIdOnJobSuccess() throws InterruptedException {
  when(bigQuery.getDataset(any(DatasetId.class))).thenReturn(dataset);
  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  when(job.waitFor(any(RetryOption.class))).thenReturn(job);
  when(job.getStatus()).thenReturn(mock(JobStatus.class));

  final BigQueryOutput bigQueryOutput = BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID);

  final StagingTableId stagingTableId = bigQueryOutput.provide(null);

  final TableId tableId = stagingTableId.publish();

  assertThat(tableId, is(TABLE_ID));
}
 

@Test(expected = RuntimeException.class)
public void shouldFailWhenJobTerminatesWithError() throws InterruptedException {
  when(bigQuery.getDataset(DATASET_ID)).thenReturn(mock(Dataset.class));

  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  when(job.waitFor(any(RetryOption.class))).thenReturn(job);
  when(job.getStatus()).thenReturn(mock(JobStatus.class));
  when(job.getStatus().getError()).thenReturn(new BigQueryError("", "", "job error"));

  BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID).provide(null).publish();
}
 

@Test(expected = RuntimeException.class)
public void shouldFailWhenJobDisappears() throws InterruptedException {
  when(bigQuery.getDataset(DATASET_ID)).thenReturn(mock(Dataset.class));

  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  when(job.waitFor(any(RetryOption.class))).thenReturn(null);

  BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID).provide(null).publish();
}
 

@Test(expected = BigQueryException.class)
public void shouldFailWhenJobTerminatesExceptionally() throws InterruptedException {
  when(bigQuery.getDataset(DATASET_ID)).thenReturn(mock(Dataset.class));

  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  doThrow(new BigQueryException(mock(IOException.class))).when(job)
      .waitFor(any(RetryOption.class));

  BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID).provide(null).publish();
}
 

/** Example of partitioning data to a list of Google Cloud Storage files. */
// [TARGET extract(String, List, JobOption...)]
// [VARIABLE "CSV"]
// [VARIABLE "gs://my_bucket/PartitionA_*.csv"]
// [VARIABLE "gs://my_bucket/PartitionB_*.csv"]
public Job extractList(String format, String gcsUrl1, String gcsUrl2) {
  // [START ]
  List<String> destinationUris = new ArrayList<>();
  destinationUris.add(gcsUrl1);
  destinationUris.add(gcsUrl2);
  Job job = table.extract(format, destinationUris);
  // Wait for the job to complete
  try {
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(3)));
    if (completedJob != null && completedJob.getStatus().getError() == null) {
      // Job completed successfully
    } else {
      // Handle error case
    }
  } catch (InterruptedException e) {
    // Handle interrupted wait
  }
  // [END ]
  return job;
}
 

/** Example loading data from a list of Google Cloud Storage files. */
// [TARGET load(FormatOptions, List, JobOption...)]
// [VARIABLE "gs://my_bucket/filename1.csv"]
// [VARIABLE "gs://my_bucket/filename2.csv"]
public Job loadList(String gcsUrl1, String gcsUrl2) {
  // [START ]
  List<String> sourceUris = new ArrayList<>();
  sourceUris.add(gcsUrl1);
  sourceUris.add(gcsUrl2);
  Job job = table.load(FormatOptions.csv(), sourceUris);
  // Wait for the job to complete
  try {
    Job completedJob =
        job.waitFor(
            RetryOption.initialRetryDelay(Duration.ofSeconds(1)),
            RetryOption.totalTimeout(Duration.ofMinutes(3)));
    if (completedJob != null && completedJob.getStatus().getError() == null) {
      // Job completed successfully
    } else {
      // Handle error case
    }
  } catch (InterruptedException e) {
    // Handle interrupted wait
  }
  // [END ]
  return job;
}