com.google.cloud.bigquery.JobStatus Java Examples

@Test
public void testLoadFile() throws IOException, ExecutionException, InterruptedException {
	ListenableFuture<Job> bigQueryJobFuture =
			bigQueryTemplate.writeDataToTable(TABLE_NAME, dataFile.getInputStream(), FormatOptions.csv());

	Job job = bigQueryJobFuture.get();
	assertThat(job.getStatus().getState()).isEqualTo(JobStatus.State.DONE);

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM test_dataset.template_test_table").build();
	TableResult result = this.bigQuery.query(queryJobConfiguration);

	assertThat(result.getTotalRows()).isEqualTo(1);
	assertThat(
			result.getValues().iterator().next().get("State").getStringValue()).isEqualTo("Alabama");
}
 

@Test
public void testLoadBytes() throws ExecutionException, InterruptedException {
	byte[] byteArray =
			"CountyId,State,County\n1001,Alabama,Autauga County\n".getBytes();
	ByteArrayInputStream byteStream = new ByteArrayInputStream(byteArray);

	ListenableFuture<Job> bigQueryJobFuture =
			bigQueryTemplate.writeDataToTable(TABLE_NAME, byteStream, FormatOptions.csv());

	Job job = bigQueryJobFuture.get();
	assertThat(job.getStatus().getState()).isEqualTo(JobStatus.State.DONE);

	QueryJobConfiguration queryJobConfiguration = QueryJobConfiguration
			.newBuilder("SELECT * FROM test_dataset.template_test_table").build();
	TableResult result = this.bigQuery.query(queryJobConfiguration);

	assertThat(result.getTotalRows()).isEqualTo(1);
	assertThat(
			result.getValues().iterator().next().get("State").getStringValue()).isEqualTo("Alabama");
}
 

@Test
public void runQuery() throws Exception {
  QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(false).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);
  when(mockJobStatus.getError()).thenReturn(null);
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);
  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);

  BigQueryDelegate delegate = new BigQueryDelegate(mockBigquery, useLegacySql);
  delegate.runQuery(queryConfig, 1000, 1000);
}
 

@Override
protected CompletableFuture<Void> close() {
  List<String> sourceUris = sourceBlobIds.stream().map(BlobIdToString::apply)
      .collect(Collectors.toList());
  boolean loadSuccess = false;
  try {
    JobStatus status = bigQuery
        .create(JobInfo.of(LoadJobConfiguration.newBuilder(tableId, sourceUris)
            .setCreateDisposition(JobInfo.CreateDisposition.CREATE_NEVER)
            .setWriteDisposition(JobInfo.WriteDisposition.WRITE_APPEND)
            .setFormatOptions(FormatOptions.json()).setIgnoreUnknownValues(true)
            .setAutodetect(false).setMaxBadRecords(0).build()))
        .waitFor().getStatus();
    if (status.getError() != null) {
      throw new BigQueryErrors(ImmutableList.of(status.getError()));
    } else if (status.getExecutionErrors() != null
        && status.getExecutionErrors().size() > 0) {
      throw new BigQueryErrors(status.getExecutionErrors());
    }
    loadSuccess = true;
    return CompletableFuture.completedFuture(null);
  } catch (InterruptedException e) {
    throw new RuntimeException(e);
  } finally {
    if (delete == Delete.always || (delete == Delete.onSuccess && loadSuccess)) {
      try {
        storage.delete(sourceBlobIds);
      } catch (RuntimeException ignore2) {
        // don't fail batch when delete throws
      }
    }
  }
}
 

/** Prepare a mock BQ response. */
@Before
public void setupMock() throws InterruptedException {
  storage = mock(Storage.class);
  bigQuery = mock(com.google.cloud.bigquery.BigQuery.class);
  job = mock(Job.class);
  jobStatus = mock(JobStatus.class);
  when(bigQuery.create(any(JobInfo.class), any())).thenReturn(job);
  when(job.waitFor(any())).thenReturn(job);
  when(job.getStatus()).thenReturn(jobStatus);
}
 

@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 = StageException.class)
public void runQueryTimeout() throws Exception {
  QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(false).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);
  when(mockJobStatus.getError()).thenReturn(null);
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);
  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);

  BigQueryDelegate delegate = new BigQueryDelegate(
      mockBigquery,
      useLegacySql,
      Clock.offset(Clock.systemDefaultZone(), Duration.ofSeconds(2))
  );

  ErrorCode code = null;
  try {
    delegate.runQuery(queryConfig, 1000, 1000);
  } catch (StageException e) {
    code = e.getErrorCode();
    throw e;
  } finally {
    assertEquals(Errors.BIGQUERY_00, code);
  }
}
 

@Test(expected = StageException.class)
public void runQueryHasErrors() throws Exception {
  QueryJobConfiguration queryRequest = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);

  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);
  when(mockJobStatus.getError()).thenReturn(new BigQueryError(
      "Some Error",
      "Some Location",
      "Some Error Message"
  ));
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);

  BigQueryDelegate delegate = new BigQueryDelegate(mockBigquery, useLegacySql);

  ErrorCode code = null;
  try {
    delegate.runQuery(queryRequest, 1000, 1000);
  } catch (StageException e) {
    code = e.getErrorCode();
    throw e;
  } finally {
    assertEquals(Errors.BIGQUERY_02, code);
  }
}
 

/** Example of reloading all fields until job status is DONE. */
// [TARGET reload(JobOption...)]
public JobStatus.State reload() throws InterruptedException {
  // [START ]
  while (!JobStatus.State.DONE.equals(job.getStatus().getState())) {
    Thread.sleep(1000L);
    job = job.reload();
  }
  // [END ]
  return job.getStatus().getState();
}
 

/** Example of reloading status field until job status is DONE. */
// [TARGET reload(JobOption...)]
public JobStatus.State reloadStatus() throws InterruptedException {
  // [START ]
  while (!JobStatus.State.DONE.equals(job.getStatus().getState())) {
    Thread.sleep(1000L);
    job = job.reload(BigQuery.JobOption.fields(BigQuery.JobField.STATUS));
  }
  // [END ]
  return job.getStatus().getState();
}
 

@Test
public void testReload() throws Exception {
  JobConfiguration jobConfig =
      QueryJobConfiguration.newBuilder(QUERY).setUseLegacySql(false).build();
  JobInfo jobInfo = JobInfo.newBuilder(jobConfig).build();
  Job job = bigquery.create(jobInfo);
  JobSnippets jobSnippets = new JobSnippets(job);
  JobStatus.State result = jobSnippets.reload();
  assertEquals(JobStatus.State.DONE, result);
}
 

@Test
public void testReloadStatus() throws Exception {
  JobConfiguration jobConfig =
      QueryJobConfiguration.newBuilder(QUERY).setUseLegacySql(false).build();
  JobInfo jobInfo = JobInfo.newBuilder(jobConfig).build();
  Job job = bigquery.create(jobInfo);
  JobSnippets jobSnippets = new JobSnippets(job);
  JobStatus.State result = jobSnippets.reloadStatus();
  assertEquals(JobStatus.State.DONE, result);
}
 

@Test
public void testProduce() throws Exception {
  File tempFile = File.createTempFile("gcp", "json");
  tempFile.deleteOnExit();

  BigQuerySourceConfig conf = new BigQuerySourceConfig();
  conf.credentials.projectId = "test";
  conf.credentials.path = tempFile.getAbsolutePath();
  conf.credentials.credentialsProvider = CredentialsProviderType.JSON_PROVIDER;
  conf.query = "SELECT * FROM [test:table]";


  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(false).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);
  when(mockJobStatus.getError()).thenReturn(null);
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  List<FieldValueList> resultSet = new ArrayList<>(1);
  resultSet.add(TestBigQueryDelegate.createTestValues());

  when(mockResult.getSchema()).thenReturn(TestBigQueryDelegate.createTestSchema());
  when(mockResult.iterateAll()).thenReturn(resultSet);
  when(mockResult.getValues()).thenReturn(resultSet);

  when(mockJob.getQueryResults()).thenReturn(mockResult);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);

  BigQuerySource bigquerySource = spy(new BigQuerySource(conf));
  doReturn(mockBigquery).when(bigquerySource).getBigQuery(any());
  doReturn(mockResult).when(bigquerySource).runQuery(any(), anyLong());

  SourceRunner runner = new SourceRunner.Builder(BigQueryDSource.class, bigquerySource)
      .addOutputLane("lane")
      .build();

  try {
    runner.runInit();

    StageRunner.Output output = runner.runProduce(null, 1000);
    List<Record> records = output.getRecords().get("lane");
    assertEquals(1, records.size());
    assertEquals(10, records.get(0).get().getValueAsListMap().size());
    assertEquals("nested string", records.get(0).get("/j/x").getValueAsString());
    assertEquals("z", records.get(0).get("/j/y/z").getValueAsString());
  } finally {
    runner.runDestroy();
  }
}