import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
@Override
void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception {
  PipelineOptions options = c.getPipelineOptions();
  BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
  String jobUuid = c.getJobId();
  final String extractDestinationDir =
      resolveTempLocation(bqOptions.getTempLocation(), "BigQueryExtractTemp", jobUuid);
  final String executingProject = bqOptions.getProject();
  JobReference jobRef =
      new JobReference()
          .setProjectId(executingProject)
          .setJobId(getExtractJobId(createJobIdToken(bqOptions.getJobName(), jobUuid)));

  Job extractJob = getBigQueryServices().getJobService(bqOptions).getJob(jobRef);

  if (extractJob != null) {
    List<ResourceId> extractFiles =
        getExtractFilePaths(extractDestinationDir, extractJob);
    if (extractFiles != null && !extractFiles.isEmpty()) {
      FileSystems.delete(
          extractFiles, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES);
    }
  }
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
private void executeQuery(
    String executingProject,
    String jobId,
    TableReference destinationTable,
    JobService jobService) throws IOException, InterruptedException {
  JobReference jobRef = new JobReference()
      .setProjectId(executingProject)
      .setJobId(jobId);

  JobConfigurationQuery queryConfig = createBasicQueryConfig()
      .setAllowLargeResults(true)
      .setCreateDisposition("CREATE_IF_NEEDED")
      .setDestinationTable(destinationTable)
      .setPriority("BATCH")
      .setWriteDisposition("WRITE_EMPTY");

  jobService.startQueryJob(jobRef, queryConfig);
  Job job = jobService.pollJob(jobRef, JOB_POLL_MAX_RETRIES);
  if (BigQueryHelpers.parseStatus(job) != Status.SUCCEEDED) {
    throw new IOException(String.format(
        "Query job %s failed, status: %s.", jobId,
        BigQueryHelpers.statusToPrettyString(job.getStatus())));
  }
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds.
 */
@Test
public void testStartLoadJobSucceeds() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

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

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob, new ApiErrorExtractor(), bigquery, sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
  expectedLogs.verifyInfo(String.format("Started BigQuery job: %s", jobRef));
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds
 * with an already exist job.
 */
@Test
public void testStartLoadJobSucceedsAlreadyExists() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

  when(response.getStatusCode()).thenReturn(409); // 409 means already exists

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob, new ApiErrorExtractor(), bigquery, sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
  expectedLogs.verifyNotLogged("Started BigQuery job");
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds with a retry.
 */
@Test
public void testStartLoadJobRetry() throws IOException, InterruptedException {
  Job testJob = new Job();
  JobReference jobRef = new JobReference();
  jobRef.setJobId("jobId");
  jobRef.setProjectId("projectId");
  testJob.setJobReference(jobRef);

  // First response is 403 rate limited, second response has valid payload.
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(403).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
      .thenReturn(toStream(testJob));

  Sleeper sleeper = new FastNanoClockAndSleeper();
  JobServiceImpl.startJob(
      testJob, new ApiErrorExtractor(), bigquery, sleeper,
      BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));

  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * 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();
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * 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();
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * 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();
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
@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();
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
@Test
public void testGetJobNotFound() throws Exception {
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(404);

  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(null, job);
  verify(response, times(1)).getStatusCode();
  verify(response, times(1)).getContent();
  verify(response, times(1)).getContentType();
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
@Override
public void startExtractJob(JobReference jobRef, JobConfigurationExtract extractConfig)
    throws InterruptedException, IOException {
  checkArgument(extractConfig.getDestinationFormat().equals("AVRO"),
      "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));
  }
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * Returns true if the provided job succeeded, false if it failed, and throws an exception if it
 * is still pending.
 */
private boolean checkJobOutcome() {
  Job job = null;
  String jobRefString =
      toJobReferenceString(new JobReference().setProjectId(projectId).setJobId(jobId));

  try {
    job = bigquery.jobs().get(projectId, jobId).execute();
  } catch (IOException e) {
    // We will throw a new exception because done==false, but first log this exception.
    logger.warning(e, e.getMessage());
  }
  // If job is not yet done, then throw an exception so that we'll return a failing HTTP status
  // code and the task will be retried.
  if (job == null || !job.getStatus().getState().equals("DONE")) {
    throw new NotModifiedException(jobRefString);
  }

  // Check if the job ended with an error.
  if (job.getStatus().getErrorResult() != null) {
    logger.severefmt("Bigquery job failed - %s - %s", jobRefString, job);
    return false;
  }
  logger.infofmt("Bigquery job succeeded - %s", jobRefString);
  return true;
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
private Job makeLoadJob(JobReference jobRef, String sourceUri, String tableId) {
  TableReference tableReference = new TableReference()
      .setProjectId(jobRef.getProjectId())
      .setDatasetId(SNAPSHOTS_DATASET)
      .setTableId(tableId);
  return new Job()
      .setJobReference(jobRef)
      .setConfiguration(new JobConfiguration()
          .setLoad(new JobConfigurationLoad()
              .setWriteDisposition(WriteDisposition.WRITE_EMPTY.toString())
              .setSourceFormat(SourceFormat.DATASTORE_BACKUP.toString())
              .setSourceUris(ImmutableList.of(sourceUri))
              .setDestinationTable(tableReference)));
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
@Test
public void testSuccess_enqueuePollTask_withChainedTask() throws Exception {
  TaskOptions chainedTask = TaskOptions.Builder
      .withUrl("/_dr/something")
      .method(Method.POST)
      .header("X-Testing", "foo")
      .param("testing", "bar");
  new BigqueryPollJobEnqueuer(ENQUEUER).enqueuePollTask(
      new JobReference().setProjectId(PROJECT_ID).setJobId(JOB_ID),
      chainedTask,
      getQueue(CHAINED_QUEUE_NAME));
  assertTasksEnqueued(BigqueryPollJobAction.QUEUE, newPollJobTaskMatcher("POST"));
  TaskStateInfo taskInfo = getOnlyElement(
      TaskQueueHelper.getQueueInfo(BigqueryPollJobAction.QUEUE).getTaskInfo());
  ByteArrayInputStream taskBodyBytes = new ByteArrayInputStream(taskInfo.getBodyAsBytes());
  TaskOptions taskOptions = (TaskOptions) new ObjectInputStream(taskBodyBytes).readObject();
  assertThat(taskOptions).isEqualTo(chainedTask);
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void startLoadJob(
    JobReference jobRef,
    JobConfigurationLoad loadConfig) throws InterruptedException, IOException {
  Job job = new Job()
      .setJobReference(jobRef)
      .setConfiguration(new JobConfiguration().setLoad(loadConfig));

  startJob(job, errorExtractor, client);
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void startExtractJob(JobReference jobRef, JobConfigurationExtract extractConfig)
    throws InterruptedException, IOException {
  Job job = new Job()
      .setJobReference(jobRef)
      .setConfiguration(
          new JobConfiguration().setExtract(extractConfig));

  startJob(job, errorExtractor, client);
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void startQueryJob(JobReference jobRef, JobConfigurationQuery queryConfig)
    throws IOException, InterruptedException {
  Job job = new Job()
      .setJobReference(jobRef)
      .setConfiguration(
          new JobConfiguration().setQuery(queryConfig));

  startJob(job, errorExtractor, client);
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
/**
 * {@inheritDoc}
 *
 * <p>Tries executing the RPC for at most {@code MAX_RPC_RETRIES} times until it succeeds.
 *
 * @throws IOException if it exceeds {@code MAX_RPC_RETRIES} attempts.
 */
@Override
public void startCopyJob(JobReference jobRef, JobConfigurationTableCopy copyConfig)
    throws IOException, InterruptedException {
  Job job = new Job()
      .setJobReference(jobRef)
      .setConfiguration(
          new JobConfiguration().setCopy(copyConfig));

  startJob(job, errorExtractor, client);
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
@Override
public Job pollJob(JobReference jobRef, int maxAttempts) throws InterruptedException {
  BackOff backoff =
      BackOffAdapter.toGcpBackOff(
          FluentBackoff.DEFAULT
              .withMaxRetries(maxAttempts)
              .withInitialBackoff(INITIAL_JOB_STATUS_POLL_BACKOFF)
              .withMaxBackoff(Duration.standardMinutes(1))
              .backoff());
  return pollJob(jobRef, Sleeper.DEFAULT, backoff);
}
 

import com.google.api.services.bigquery.model.JobReference; //导入依赖的package包/类
private List<ResourceId> executeExtract(
    String jobId, TableReference table, JobService jobService, String executingProject,
    String extractDestinationDir)
        throws InterruptedException, IOException {
  JobReference jobRef = new JobReference()
      .setProjectId(executingProject)
      .setJobId(jobId);

  String destinationUri = BigQueryIO.getExtractDestinationUri(extractDestinationDir);
  JobConfigurationExtract extract = new JobConfigurationExtract()
      .setSourceTable(table)
      .setDestinationFormat("AVRO")
      .setDestinationUris(ImmutableList.of(destinationUri));

  LOG.info("Starting BigQuery extract job: {}", jobId);
  jobService.startExtractJob(jobRef, extract);
  Job extractJob =
      jobService.pollJob(jobRef, JOB_POLL_MAX_RETRIES);
  if (BigQueryHelpers.parseStatus(extractJob) != Status.SUCCEEDED) {
    throw new IOException(String.format(
        "Extract job %s failed, status: %s.",
        extractJob.getJobReference().getJobId(),
        BigQueryHelpers.statusToPrettyString(extractJob.getStatus())));
  }

  LOG.info("BigQuery extract job completed: {}", jobId);

  return BigQueryIO.getExtractFilePaths(extractDestinationDir, extractJob);
}