Java Code Examples for com.google.api.services.bigquery.model.Job#setStatus()

@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));
  }
}
 

/** 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();
}
 

/** 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} 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();
}
 

@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 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));
  }
}
 

@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 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));
  }
}
 

/**
 * 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);
}
 

@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);
}
 

@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);
}
 

/**
 * 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);
}