com.google.api.services.bigquery.model.JobStatus Java Examples

/**
 * Checks completed job for errors.
 *
 * @throws BigqueryJobFailureException
 */
private static Job checkJob(Job job) {
  verify(job.getStatus() != null);
  JobStatus jobStatus = job.getStatus();
  if (jobStatus.getErrorResult() != null) {
    throw BigqueryJobFailureException.create(jobStatus);
  } else {
    logger.atInfo().log(summarizeCompletedJob(job));
    if (jobStatus.getErrors() != null) {
      for (ErrorProto error : jobStatus.getErrors()) {
        logger.atWarning().log("%s: %s", error.getReason(), error.getMessage());
      }
    }
    return job;
  }
}
 

@Override
public void startExtractJob(JobReference jobRef, JobConfigurationExtract extractConfig)
    throws IOException {
  checkArgument(
      "AVRO".equals(extractConfig.getDestinationFormat()), "Only extract to AVRO is supported");
  synchronized (allJobs) {
    verifyUniqueJobId(jobRef.getJobId());
    ++numExtractJobCalls;

    Job job = new Job();
    job.setJobReference(jobRef);
    job.setConfiguration(new JobConfiguration().setExtract(extractConfig));
    job.setKind(" bigquery#job");
    job.setStatus(new JobStatus().setState("PENDING"));
    allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
  }
}
 

@Test
public void testGetJobSucceeds() throws Exception {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus());

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.getJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

@Override
public Job pollJob(JobReference jobRef, int maxAttempts) throws InterruptedException {
  BackOff backoff =
      BackOffAdapter.toGcpBackOff(
          FluentBackoff.DEFAULT
              .withMaxRetries(maxAttempts)
              .withInitialBackoff(Duration.millis(10))
              .withMaxBackoff(Duration.standardSeconds(1))
              .backoff());
  Sleeper sleeper = Sleeper.DEFAULT;
  try {
    do {
      Job job = getJob(jobRef);
      if (job != null) {
        JobStatus status = job.getStatus();
        if (status != null
            && ("DONE".equals(status.getState()) || "FAILED".equals(status.getState()))) {
          return job;
        }
      }
    } while (BackOffUtils.next(sleeper, backoff));
  } catch (IOException e) {
    return null;
  }
  return null;
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} returns UNKNOWN. */
@Test
public void testPollJobUnknown() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus());

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.STOP_BACKOFF);

  assertEquals(null, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} fails. */
@Test
public void testPollJobFailed() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus().setState("DONE").setErrorResult(new ErrorProto()));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

/** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} succeeds. */
@Test
public void testPollJobSucceeds() throws IOException, InterruptedException {
  Job testJob = new Job();
  testJob.setStatus(new JobStatus().setState("DONE"));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent()).thenReturn(toStream(testJob));

  BigQueryServicesImpl.JobServiceImpl jobService =
      new BigQueryServicesImpl.JobServiceImpl(bigquery);
  JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
  Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);

  assertEquals(testJob, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

private JobStatus runExtractJob(Job job, JobConfigurationExtract extract)
    throws InterruptedException, IOException {
  TableReference sourceTable = extract.getSourceTable();

  List<TableRow> rows =
      datasetService.getAllRows(
          sourceTable.getProjectId(), sourceTable.getDatasetId(), sourceTable.getTableId());
  TableSchema schema = datasetService.getTable(sourceTable).getSchema();
  List<Long> destinationFileCounts = Lists.newArrayList();
  for (String destination : extract.getDestinationUris()) {
    destinationFileCounts.add(writeRows(sourceTable.getTableId(), rows, schema, destination));
  }
  job.setStatistics(
      new JobStatistics()
          .setExtract(new JobStatistics4().setDestinationUriFileCounts(destinationFileCounts)));
  return new JobStatus().setState("DONE");
}
 

@Test
public void testJobFailed() throws Exception {
  when(bigqueryJobsGet.execute()).thenReturn(new Job().setStatus(
      new JobStatus()
          .setState("DONE")
          .setErrorResult(new ErrorProto().setMessage("Job failed"))));
  action.run();
  assertLogMessage(
      logHandler, SEVERE, String.format("Bigquery job failed - %s:%s", PROJECT_ID, JOB_ID));
  assertNoTasksEnqueued(CHAINED_QUEUE_NAME);
}
 

@Test
public void testSuccess_chainedPayloadAndJobSucceeded_enqueuesChainedTask() throws Exception {
  when(bigqueryJobsGet.execute()).thenReturn(
      new Job().setStatus(new JobStatus().setState("DONE")));

  TaskOptions chainedTask =
      TaskOptions.Builder.withUrl("/_dr/something")
          .method(Method.POST)
          .header("X-Testing", "foo")
          .param("testing", "bar")
          .taskName("my_task_name");
  ByteArrayOutputStream bytes = new ByteArrayOutputStream();
  new ObjectOutputStream(bytes).writeObject(chainedTask);
  action.payload = bytes.toByteArray();

  action.run();
  assertLogMessage(
      logHandler, INFO, String.format("Bigquery job succeeded - %s:%s", PROJECT_ID, JOB_ID));
  assertLogMessage(
      logHandler,
      INFO,
      "Added chained task my_task_name for /_dr/something to queue " + CHAINED_QUEUE_NAME);
  assertTasksEnqueued(
      CHAINED_QUEUE_NAME,
      new TaskMatcher()
          .url("/_dr/something")
          .method("POST")
          .header("X-Testing", "foo")
          .param("testing", "bar")
          .taskName("my_task_name"));
}
 

@Test
public void testSuccess_jobCompletedSuccessfully() throws Exception {
  when(bigqueryJobsGet.execute()).thenReturn(
      new Job().setStatus(new JobStatus().setState("DONE")));
  action.run();
  assertLogMessage(
      logHandler, INFO, String.format("Bigquery job succeeded - %s:%s", PROJECT_ID, JOB_ID));
}
 

public BigqueryJobFailureException(
    String message,
    @Nullable Throwable cause,
    @Nullable JobStatus jobStatus,
    @Nullable GoogleJsonError jsonError) {
  super(message, cause);
  this.jobStatus = jobStatus;
  this.jsonError = jsonError;
}
 

@Test
public void testFailure_badChainedTaskPayload() throws Exception {
  when(bigqueryJobsGet.execute()).thenReturn(
      new Job().setStatus(new JobStatus().setState("DONE")));
  action.payload = "payload".getBytes(UTF_8);
  BadRequestException thrown = assertThrows(BadRequestException.class, action::run);
  assertThat(thrown).hasMessageThat().contains("Cannot deserialize task from payload");
}
 

private JobStatus runQueryJob(JobConfigurationQuery query)
    throws IOException, InterruptedException {
  KV<Table, List<TableRow>> result = FakeBigQueryServices.decodeQueryResult(query.getQuery());
  datasetService.createTable(result.getKey().setTableReference(query.getDestinationTable()));
  datasetService.insertAll(query.getDestinationTable(), result.getValue(), null);
  return new JobStatus().setState("DONE");
}
 

/**
 * Mocks result of BigQuery for polling for job completion.
 *
 * @throws IOException on IOError.
 */
@Before
public void setUp() throws IOException {

  jobReference = new JobReference().setJobId("test-job-id").setLocation("test-job-location");

  // Create the unfinished job result.
  notDoneJob = new Job();
  notDoneJobStatus = new JobStatus();
  notDoneJobStatus.setState("NOT DONE");
  notDoneJobStatus.setErrorResult(null);
  notDoneJob.setStatus(notDoneJobStatus);
  notDoneJob.setJobReference(jobReference);

  // Create the finished job result.
  job = new Job();
  jobStatus = new JobStatus();
  jobStatus.setState("DONE");
  jobStatus.setErrorResult(null);
  job.setStatus(jobStatus);
  job.setJobReference(jobReference);

  // Mock BigQuery.
  mockBigQuery = mock(Bigquery.class);
  mockBigQueryJobs = mock(Bigquery.Jobs.class);
  mockJobsGet = mock(Bigquery.Jobs.Get.class);
  when(mockBigQuery.jobs()).thenReturn(mockBigQueryJobs);
  when(mockBigQueryJobs.get(projectId, jobReference.getJobId()))
      .thenReturn(mockJobsGet)
      .thenReturn(mockJobsGet);
  when(mockJobsGet.setLocation(any(String.class))).thenReturn(mockJobsGet);
  when(mockJobsGet.execute()).thenReturn(job);

  // Constructor coverage
  new BigQueryUtils();

  // Mock Progressable.
  mockProgressable = mock(Progressable.class);
}
 

private JobStatus runJob(Job job) throws InterruptedException, IOException {
  if (job.getConfiguration().getLoad() != null) {
    return runLoadJob(job.getJobReference(), job.getConfiguration().getLoad());
  } else if (job.getConfiguration().getCopy() != null) {
    return runCopyJob(job.getConfiguration().getCopy());
  } else if (job.getConfiguration().getExtract() != null) {
    return runExtractJob(job, job.getConfiguration().getExtract());
  } else if (job.getConfiguration().getQuery() != null) {
    return runQueryJob(job.getConfiguration().getQuery());
  }
  return new JobStatus().setState("DONE");
}
 

@Override
public void startCopyJob(JobReference jobRef, JobConfigurationTableCopy copyConfig)
    throws IOException {
  synchronized (allJobs) {
    verifyUniqueJobId(jobRef.getJobId());
    Job job = new Job();
    job.setJobReference(jobRef);
    job.setConfiguration(new JobConfiguration().setCopy(copyConfig));
    job.setKind(" bigquery#job");
    job.setStatus(new JobStatus().setState("PENDING"));
    allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
  }
}
 

@Override
public void startQueryJob(JobReference jobRef, JobConfigurationQuery query) {
  synchronized (allJobs) {
    Job job = new Job();
    job.setJobReference(jobRef);
    job.setConfiguration(new JobConfiguration().setQuery(query));
    job.setKind(" bigquery#job");
    job.setStatus(new JobStatus().setState("PENDING"));
    allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
  }
}
 

@Override
public void startLoadJob(JobReference jobRef, JobConfigurationLoad loadConfig)
    throws IOException {
  synchronized (allJobs) {
    verifyUniqueJobId(jobRef.getJobId());
    Job job = new Job();
    job.setJobReference(jobRef);
    job.setConfiguration(new JobConfiguration().setLoad(loadConfig));
    job.setKind(" bigquery#job");
    job.setStatus(new JobStatus().setState("PENDING"));

    // Copy the files to a new location for import, as the temporary files will be deleted by
    // the caller.
    if (loadConfig.getSourceUris().size() > 0) {
      ImmutableList.Builder<ResourceId> sourceFiles = ImmutableList.builder();
      ImmutableList.Builder<ResourceId> loadFiles = ImmutableList.builder();
      for (String filename : loadConfig.getSourceUris()) {
        sourceFiles.add(FileSystems.matchNewResource(filename, false /* isDirectory */));
        loadFiles.add(
            FileSystems.matchNewResource(
                filename + ThreadLocalRandom.current().nextInt(), false /* isDirectory */));
      }

      FileSystems.copy(sourceFiles.build(), loadFiles.build());
      filesForLoadJobs.put(jobRef.getProjectId(), jobRef.getJobId(), loadFiles.build());
    }

    allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
  }
}
 

@VisibleForTesting
Job pollJob(JobReference jobRef, Sleeper sleeper, BackOff backoff) throws InterruptedException {
  do {
    try {
      Job job =
          client
              .jobs()
              .get(jobRef.getProjectId(), jobRef.getJobId())
              .setLocation(jobRef.getLocation())
              .execute();
      if (job == null) {
        LOG.info("Still waiting for BigQuery job {} to start", jobRef);
        continue;
      }
      JobStatus status = job.getStatus();
      if (status == null) {
        LOG.info("Still waiting for BigQuery job {} to enter pending state", jobRef);
        continue;
      }
      if ("DONE".equals(status.getState())) {
        LOG.info("BigQuery job {} completed in state DONE", jobRef);
        return job;
      }
      // The job is not DONE, wait longer and retry.
      LOG.info(
          "Still waiting for BigQuery job {}, currently in status {}\n{}",
          jobRef.getJobId(),
          status,
          formatBqStatusCommand(jobRef.getProjectId(), jobRef.getJobId()));
    } catch (IOException e) {
      // ignore and retry
      LOG.info("Ignore the error and retry polling job status.", e);
    }
  } while (nextBackOff(sleeper, backoff));
  LOG.warn("Unable to poll job status: {}, aborting after reached max .", jobRef.getJobId());
  return null;
}
 

static Status parseStatus(@Nullable Job job) {
  if (job == null) {
    return Status.UNKNOWN;
  }
  JobStatus status = job.getStatus();
  if (status.getErrorResult() != null) {
    return Status.FAILED;
  } else if (status.getErrors() != null && !status.getErrors().isEmpty()) {
    return Status.FAILED;
  } else {
    return Status.SUCCEEDED;
  }
}
 

@Before
public void setUp() throws IOException {
  MockitoAnnotations.initMocks(this);
  LoggerConfig.getConfig(GsonBigQueryInputFormat.class).setLevel(Level.FINE);

  // Create fake job reference.
  JobReference fakeJobReference = new JobReference().setProjectId(jobProjectId).setJobId(jobId);

  // Create the job result.
  jobStatus = new JobStatus();
  jobStatus.setState("DONE");
  jobStatus.setErrorResult(null);

  jobHandle = new Job();
  jobHandle.setStatus(jobStatus);
  jobHandle.setJobReference(fakeJobReference);

  // Mocks for Bigquery jobs.
  when(mockBigquery.jobs()).thenReturn(mockBigqueryJobs);

  // Mock getting Bigquery job.
  when(mockBigqueryJobs.get(any(String.class), any(String.class)))
      .thenReturn(mockBigqueryJobsGet);
  when(mockBigqueryJobsGet.setLocation(any(String.class))).thenReturn(mockBigqueryJobsGet);

  // Mock inserting Bigquery job.
  when(mockBigqueryJobs.insert(any(String.class), any(Job.class)))
      .thenReturn(mockBigqueryJobsInsert);

  // Fake table.
  fakeTableSchema = new TableSchema();
  fakeTable = new Table().setSchema(fakeTableSchema).setLocation("test_location");

  // Mocks for Bigquery tables.
  when(mockBigquery.tables()).thenReturn(mockBigqueryTables);
  when(mockBigqueryTables.get(any(String.class), any(String.class), any(String.class)))
      .thenReturn(mockBigqueryTablesGet);

  Datasets datasets = Mockito.mock(Datasets.class);
  Datasets.Get datasetsGet = Mockito.mock(Datasets.Get.class);
  Dataset dataset = new Dataset().setLocation("test_location");
  when(mockBigquery.datasets()).thenReturn(datasets);
  when(datasets.get(any(String.class), any(String.class))).thenReturn(datasetsGet);
  when(datasetsGet.execute()).thenReturn(dataset);

  // Create table reference.
  tableRef = new TableReference();
  tableRef.setProjectId(projectId);
  tableRef.setDatasetId(datasetId);
  tableRef.setTableId(tableId);

  helper = new BigQueryHelper(mockBigquery);
  helper.setErrorExtractor(mockErrorExtractor);
}
 

@Test
public void testJobPending() throws Exception {
  when(bigqueryJobsGet.execute()).thenReturn(
      new Job().setStatus(new JobStatus().setState("PENDING")));
  assertThrows(NotModifiedException.class, action::run);
}
 

static RetryJobIdResult getRetryJobId(
    RetryJobId currentJobId, SerializableFunction<RetryJobId, Job> lookupJob) {
  for (int retryIndex = currentJobId.getRetryIndex(); ; retryIndex++) {
    RetryJobId jobId = new RetryJobId(currentJobId.getJobIdPrefix(), retryIndex);
    try {
      Job loadJob = lookupJob.apply(jobId);
      if (loadJob == null) {
        LOG.info("job id {} not found, so retrying with that id", jobId);
        // This either means that the original job was never properly issued (on the first
        // iteration of the loop) or that we've found a retry id that has not been used yet. Try
        // again with this job id.
        return new RetryJobIdResult(jobId, true);
      }
      JobStatus jobStatus = loadJob.getStatus();
      if (jobStatus == null) {
        LOG.info("job status for {} not found, so retrying with that job id", jobId);
        return new RetryJobIdResult(jobId, true);
      }
      if ("PENDING".equals(jobStatus.getState()) || "RUNNING".equals(jobStatus.getState())) {
        // The job id has been issued and is currently pending. This can happen after receiving
        // an error from the load or copy job creation (e.g. that error might come because the
        // job already exists). Return to the caller which job id is pending (it might not be the
        // one passed in) so the caller can then wait for this job to finish.
        LOG.info("job {} in pending or running state, so continuing with that job id", jobId);
        return new RetryJobIdResult(jobId, false);
      }
      if (jobStatus.getErrorResult() == null
          && (jobStatus.getErrors() == null || jobStatus.getErrors().isEmpty())) {
        // Import succeeded. No retry needed.
        LOG.info("job {} succeeded, so not retrying ", jobId);
        return new RetryJobIdResult(jobId, false);
      }
      // This job has failed, so we assume the data cannot enter BigQuery. We will check the next
      // job in the sequence (with the same unique prefix) to see if is either pending/succeeded
      // or can be used to generate a retry job.
      LOG.info("job {} is failed. Checking the next job id", jobId);
    } catch (RuntimeException e) {
      LOG.info("caught exception while querying job {}", jobId);
      return new RetryJobIdResult(jobId, true);
    }
  }
}
 

/**
 * Creates an in-memory GHFS.
 *
 * @throws IOException on IOError.
 */
@Before
public void setUp()
    throws IOException {
  MockitoAnnotations.initMocks(this);
  LoggerConfig.getConfig(GsonBigQueryInputFormat.class).setLevel(Level.FINE);

  // Set the Hadoop job configuration.
  config = new JobConf(InMemoryGoogleHadoopFileSystem.getSampleConfiguration());
  config.set(BigQueryConfiguration.PROJECT_ID.getKey(), jobProjectId);
  config.set(BigQueryConfiguration.INPUT_PROJECT_ID.getKey(), dataProjectId);
  config.set(BigQueryConfiguration.INPUT_DATASET_ID.getKey(), intermediateDataset);
  config.set(BigQueryConfiguration.INPUT_TABLE_ID.getKey(), intermediateTable);
  config.set(BigQueryConfiguration.TEMP_GCS_PATH.getKey(), "gs://test_bucket/other_path");
  config.setClass(
      INPUT_FORMAT_CLASS.getKey(),
      GsonBigQueryInputFormat.class,
      AbstractBigQueryInputFormat.class);
  config.setBoolean(BigQueryConfiguration.DELETE_EXPORT_FILES_FROM_GCS.getKey(), true);

  CredentialConfigurationUtil.addTestConfigurationSettings(config);

  // Create a GoogleHadoopFileSystem to use to initialize and write to
  // the in-memory GcsFs.
  ghfs = new InMemoryGoogleHadoopFileSystem();

  JobReference fakeJobReference =
      new JobReference()
          .setProjectId(jobProjectId)
          .setJobId("bigquery-job-1234")
          .setLocation("test-job-location");

  // Create the job result.
  jobStatus = new JobStatus();
  jobStatus.setState("DONE");
  jobStatus.setErrorResult(null);

  jobHandle = new Job();
  jobHandle.setStatus(jobStatus);
  jobHandle.setJobReference(fakeJobReference);

  // Create table reference.
  tableRef = new TableReference();
  tableRef.setProjectId(dataProjectId);
  tableRef.setDatasetId("test_dataset");
  tableRef.setTableId("test_table");

  table = new Table().setTableReference(tableRef).setLocation("test_location");

  when(mockBigQueryHelper.getRawBigquery())
      .thenReturn(mockBigquery);

  // Mocks for Bigquery jobs.
  when(mockBigquery.jobs())
      .thenReturn(mockBigqueryJobs);

  // Mock getting Bigquery job.
  when(mockBigqueryJobs.get(any(String.class), any(String.class)))
      .thenReturn(mockBigqueryJobsGet);
  when(mockBigqueryJobsGet.setLocation(any(String.class))).thenReturn(mockBigqueryJobsGet);
  when(mockBigqueryJobsGet.execute())
      .thenReturn(jobHandle);

  // Mock inserting Bigquery job.
  when(mockBigqueryJobs.insert(any(String.class), any(Job.class)))
      .thenReturn(mockBigqueryJobsInsert);
  when(mockBigqueryJobsInsert.execute())
      .thenReturn(jobHandle);

  // Mocks for Bigquery tables.
  when(mockBigquery.tables())
      .thenReturn(mockBigqueryTables);

  // Mocks for getting Bigquery table.
  when(mockBigqueryTables.get(any(String.class), any(String.class), any(String.class)))
      .thenReturn(mockBigqueryTablesGet);
  when(mockBigqueryTablesGet.execute())
      .thenReturn(table);

  when(mockBigQueryHelper.getTable(any(TableReference.class)))
      .thenReturn(table);

  when(mockBigQueryHelper.createJobReference(
          any(String.class), any(String.class), any(String.class)))
      .thenReturn(fakeJobReference);
  when(mockBigQueryHelper.insertJobOrFetchDuplicate(any(String.class), any(Job.class)))
      .thenReturn(jobHandle);
}
 

/** Create an error from a failed job. */
public static BigqueryJobFailureException create(JobStatus jobStatus) {
  checkArgument(jobStatus.getErrorResult() != null, "this job didn't fail!");
  return new BigqueryJobFailureException(
      describeError(jobStatus.getErrorResult()), null, jobStatus, null);
}
 

@Test
public void testPendingJobManager() throws Exception {
  PendingJobManager jobManager =
      new PendingJobManager(
          BackOffAdapter.toGcpBackOff(
              FluentBackoff.DEFAULT
                  .withMaxRetries(Integer.MAX_VALUE)
                  .withInitialBackoff(Duration.millis(10))
                  .withMaxBackoff(Duration.millis(10))
                  .backoff()));

  Set<String> succeeded = Sets.newHashSet();
  for (int i = 0; i < 5; i++) {
    Job currentJob = new Job();
    currentJob.setKind(" bigquery#job");
    PendingJob pendingJob =
        new PendingJob(
            retryId -> {
              if (new Random().nextInt(2) == 0) {
                throw new RuntimeException("Failing to start.");
              }
              currentJob.setJobReference(
                  new JobReference()
                      .setProjectId("")
                      .setLocation("")
                      .setJobId(retryId.getJobId()));
              return null;
            },
            retryId -> {
              if (retryId.getRetryIndex() < 5) {
                currentJob.setStatus(new JobStatus().setErrorResult(new ErrorProto()));
              } else {
                currentJob.setStatus(new JobStatus().setErrorResult(null));
              }
              return currentJob;
            },
            retryId -> {
              if (retryId.getJobId().equals(currentJob.getJobReference().getJobId())) {
                return currentJob;
              } else {
                return null;
              }
            },
            100,
            "JOB_" + i);
    jobManager.addPendingJob(
        pendingJob,
        j -> {
          succeeded.add(j.currentJobId.getJobId());
          return null;
        });
  }

  jobManager.waitForDone();
  Set<String> expectedJobs =
      ImmutableSet.of("JOB_0-5", "JOB_1-5", "JOB_2-5", "JOB_3-5", "JOB_4-5");
  assertEquals(expectedJobs, succeeded);
}
 

private JobStatus runCopyJob(JobConfigurationTableCopy copy)
    throws InterruptedException, IOException {
  List<TableReference> sources = copy.getSourceTables();
  TableReference destination = copy.getDestinationTable();
  WriteDisposition writeDisposition = WriteDisposition.valueOf(copy.getWriteDisposition());
  CreateDisposition createDisposition = CreateDisposition.valueOf(copy.getCreateDisposition());
  Table existingTable = datasetService.getTable(destination);
  if (!validateDispositions(existingTable, createDisposition, writeDisposition)) {
    return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto());
  }
  TimePartitioning partitioning = null;
  Clustering clustering = null;
  TableSchema schema = null;
  boolean first = true;
  List<TableRow> allRows = Lists.newArrayList();
  for (TableReference source : sources) {
    Table table = checkNotNull(datasetService.getTable(source));
    if (!first) {
      if (!Objects.equals(partitioning, table.getTimePartitioning())) {
        return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto());
      }
      if (!Objects.equals(clustering, table.getClustering())) {
        return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto());
      }
      if (!Objects.equals(schema, table.getSchema())) {
        return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto());
      }
    }
    partitioning = table.getTimePartitioning();
    clustering = table.getClustering();
    schema = table.getSchema();
    first = false;
    allRows.addAll(
        datasetService.getAllRows(
            source.getProjectId(), source.getDatasetId(), source.getTableId()));
  }
  datasetService.createTable(
      new Table()
          .setTableReference(destination)
          .setSchema(schema)
          .setTimePartitioning(partitioning)
          .setClustering(clustering)
          .setEncryptionConfiguration(copy.getDestinationEncryptionConfiguration()));
  datasetService.insertAll(destination, allRows, null);
  return new JobStatus().setState("DONE");
}
 

private JobStatus runLoadJob(JobReference jobRef, JobConfigurationLoad load)
    throws InterruptedException, IOException {
  TableReference destination = load.getDestinationTable();
  TableSchema schema = load.getSchema();
  checkArgument(schema != null, "No schema specified");
  List<ResourceId> sourceFiles = filesForLoadJobs.get(jobRef.getProjectId(), jobRef.getJobId());
  WriteDisposition writeDisposition = WriteDisposition.valueOf(load.getWriteDisposition());
  CreateDisposition createDisposition = CreateDisposition.valueOf(load.getCreateDisposition());

  Table existingTable = datasetService.getTable(destination);
  if (!validateDispositions(existingTable, createDisposition, writeDisposition)) {
    return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto());
  }
  if (existingTable == null) {
    TableReference strippedDestination =
        destination
            .clone()
            .setTableId(BigQueryHelpers.stripPartitionDecorator(destination.getTableId()));
    existingTable = new Table().setTableReference(strippedDestination).setSchema(schema);
    if (load.getTimePartitioning() != null) {
      existingTable = existingTable.setTimePartitioning(load.getTimePartitioning());
    }
    if (load.getClustering() != null) {
      existingTable = existingTable.setClustering(load.getClustering());
    }
    datasetService.createTable(existingTable);
  }

  List<TableRow> rows = Lists.newArrayList();
  for (ResourceId filename : sourceFiles) {
    if (load.getSourceFormat().equals("NEWLINE_DELIMITED_JSON")) {
      rows.addAll(readJsonTableRows(filename.toString()));
    } else if (load.getSourceFormat().equals("AVRO")) {
      rows.addAll(readAvroTableRows(filename.toString(), schema));
    }
  }

  datasetService.insertAll(destination, rows, null);
  FileSystems.delete(sourceFiles);
  return new JobStatus().setState("DONE");
}
 

static String statusToPrettyString(@Nullable JobStatus status) throws IOException {
  return status == null ? "Unknown status: null." : status.toPrettyString();
}