BigQueryUtils { public static void waitForJobCompletion( Bigquery bigquery, String projectId, JobReference jobReference, Progressable progressable) throws IOException, InterruptedException { Sleeper sleeper = Sleeper.DEFAULT; BackOff pollBackOff = new ExponentialBackOff.Builder() .setMaxIntervalMillis(POLL_WAIT_INTERVAL_MAX_MILLIS) .setInitialIntervalMillis(POLL_WAIT_INITIAL_MILLIS) .setMaxElapsedTimeMillis(POLL_WAIT_MAX_ELAPSED_MILLIS) .build(); long startTime = System.currentTimeMillis(); long elapsedTime = 0; boolean notDone = true; while (notDone) { BackOff operationBackOff = new ExponentialBackOff(); Get get = bigquery .jobs() .get(projectId, jobReference.getJobId()) .setLocation(jobReference.getLocation()); Job pollJob = ResilientOperation.retry( ResilientOperation.getGoogleRequestCallable(get), operationBackOff, RetryDeterminer.RATE_LIMIT_ERRORS, IOException.class, sleeper); elapsedTime = System.currentTimeMillis() - startTime; logger.atFine().log( "Job status (%s ms) %s: %s", elapsedTime, jobReference.getJobId(), pollJob.getStatus().getState()); if (pollJob.getStatus().getState().equals("DONE")) { notDone = false; if (pollJob.getStatus().getErrorResult() != null) { throw new IOException( String.format( "Job %s failed: %s. Errors: %s", jobReference.getJobId(), pollJob.getStatus().getErrorResult(), pollJob.getStatus().getErrors())); } } else { long millisToWait = pollBackOff.nextBackOffMillis(); if (millisToWait == BackOff.STOP) { throw new IOException( String.format( "Job %s failed to complete after %s millis.", jobReference.getJobId(), elapsedTime)); } Thread.sleep(millisToWait); progressable.progress(); } } }  static void waitForJobCompletion(
      Bigquery bigquery, String projectId, JobReference jobReference, Progressable progressable); static List<TableFieldSchema> getSchemaFromString(String fields); static final int POLL_WAIT_INITIAL_MILLIS; static final int POLL_WAIT_INTERVAL_MAX_MILLIS; static final int POLL_WAIT_MAX_ELAPSED_MILLIS; }
@Test public void testWaitForJobCompletionError() throws InterruptedException, IOException { when(mockJobsGet.execute()).thenReturn(job); jobStatus.setErrorResult(new ErrorProto()); IOException e = assertThrows( IOException.class, () -> BigQueryUtils.waitForJobCompletion( mockBigQuery, projectId, jobReference, mockProgressable)); assertThat(e).hasMessageThat().contains(jobReference.getJobId()); }
@Test public void testWaitForJobCompletion() throws Exception { when(mockJobsGet.execute()).thenReturn(job); BigQueryUtils.waitForJobCompletion(mockBigQuery, projectId, jobReference, mockProgressable); verify(mockBigQuery).jobs(); verify(mockBigQueryJobs).get(projectId, jobReference.getJobId()); verify(mockJobsGet).execute(); verify(mockProgressable, never()).progress(); }
@Test public void testWaitForJobCompletionChange() throws Exception { when(mockJobsGet.execute()).thenReturn(notDoneJob).thenReturn(job); BigQueryUtils.waitForJobCompletion(mockBigQuery, projectId, jobReference, mockProgressable); verify(mockBigQuery, times(2)).jobs(); verify(mockBigQueryJobs, times(2)).get(projectId, jobReference.getJobId()); verify(mockJobsGet, times(2)).execute(); verify(mockProgressable, atLeastOnce()).progress(); }