com.google.cloud.bigquery.JobInfo Java Examples

@Override
void run(BigQuery bigquery, JobInfo job) throws Exception {
  System.out.println("Creating job");
  Job startedJob = bigquery.create(job);
  while (!startedJob.isDone()) {
    System.out.println("Waiting for job " + startedJob.getJobId().getJob() + " to complete");
    Thread.sleep(1000L);
  }
  startedJob = startedJob.reload();
  if (startedJob.getStatus().getError() == null) {
    System.out.println("Job " + startedJob.getJobId().getJob() + " succeeded");
  } else {
    System.out.println("Job " + startedJob.getJobId().getJob() + " failed");
    System.out.println("Error: " + startedJob.getStatus().getError());
  }
}
 

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

/** Example of loading a parquet file from GCS to a table. */
public void loadTableGcsParquet(String datasetName) throws InterruptedException {
  // [START bigquery_load_table_gcs_parquet]
  String sourceUri = "gs://cloud-samples-data/bigquery/us-states/us-states.parquet";
  TableId tableId = TableId.of(datasetName, "us_states");
  LoadJobConfiguration configuration =
      LoadJobConfiguration.builder(tableId, sourceUri)
          .setFormatOptions(FormatOptions.parquet())
          .build();
  // Load the table
  Job loadJob = bigquery.create(JobInfo.of(configuration));
  loadJob = loadJob.waitFor();
  // Check the table
  StandardTableDefinition destinationTable = bigquery.getTable(tableId).getDefinition();
  System.out.println("State: " + loadJob.getStatus().getState());
  System.out.printf("Loaded %d rows.\n", destinationTable.getNumRows());
  // [END bigquery_load_table_gcs_parquet]
}
 

@Test
public void shouldRunExtractJobInTestMode() throws Exception {
  final TableId srcTable = TableId.of("foo", "bar", "baz");

  final String destinationUri = "gs://foo/bar";

  final Task<String> task = Task.named("task")
      .ofType(String.class)
      .operator(BigQueryOperator.create())
      .process(bq -> bq.job(
          JobInfo.of(ExtractJobConfiguration.of(srcTable, destinationUri)))
          .success(response -> destinationUri));

  try (TestScope scope = FloTesting.scope()) {

    final String result = FloRunner.runTask(task).future()
        .get(30, SECONDS);

    assertThat(result, is(destinationUri));
  }
}
 

@Test
public void shouldRunLoadJobInTestMode() throws Exception {
  final TableId dstTable = TableId.of("foo", "bar", "baz");
  final String srcUri = "gs://foo/bar";

  final Task<TableId> task = Task.named("task")
      .ofType(TableId.class)
      .operator(BigQueryOperator.create())
      .process(bq -> bq.job(
          JobInfo.of(LoadJobConfiguration.of(dstTable, srcUri)))
          .success(response -> dstTable));

  try (TestScope scope = FloTesting.scope()) {

    final TableId result = FloRunner.runTask(task).future()
        .get(30, SECONDS);

    assertThat(result, is(dstTable));
  }
}
 

@Test
public void shouldRunCopyJobInTestMode() throws Exception {
  final TableId srcTable = TableId.of("foo", "bar", "src");
  final TableId dstTable = TableId.of("foo", "bar", "dst");

  final Task<TableId> task = Task.named("task")
      .ofType(TableId.class)
      .operator(BigQueryOperator.create())
      .process(bq -> bq.job(
          JobInfo.of(CopyJobConfiguration.of(dstTable, srcTable)))
          .success(response -> dstTable));

  try (TestScope scope = FloTesting.scope()) {

    final TableId result = FloRunner.runTask(task).future()
        .get(30, SECONDS);

    assertThat(result, is(dstTable));
  }
}
 

@Test
public void shouldRunQueryJobInTestMode() throws Exception {
  final TableId table = TableId.of("foo", "bar", "baz");

  final Task<TableId> task = Task.named("task")
      .ofType(TableId.class)
      .output(BigQueryOutput.create(table))
      .operator(BigQueryOperator.create())
      .process((stagingTable, bq) -> bq.job(
          JobInfo.of(QueryJobConfiguration.newBuilder("SELECT foo FROM input")
              .setDestinationTable(stagingTable.tableId())
              .build()))
          .success(response -> stagingTable.publish()));

  try (TestScope scope = FloTesting.scope()) {

    final TableId result = FloRunner.runTask(task).future()
        .get(30, SECONDS);

    assertThat(result, is(table));
    assertThat(BigQueryMocking.mock().tablePublished(table), is(true));
    assertThat(BigQueryMocking.mock().tableExists(table), is(true));
  }
}
 

/** Example of copying multiple tables to a destination. */
public void copyTables(String datasetId, String destinationTableId) throws InterruptedException {
  generateTableWithDdl(datasetId, "table1");
  generateTableWithDdl(datasetId, "table2");

  // [START bigquery_copy_table_multiple_source]
  TableId destinationTable = TableId.of(datasetId, destinationTableId);
  CopyJobConfiguration configuration =
      CopyJobConfiguration.newBuilder(
              destinationTable,
              Arrays.asList(TableId.of(datasetId, "table1"), TableId.of(datasetId, "table2")))
          .build();

  // Copy the tables.
  Job job = bigquery.create(JobInfo.of(configuration));
  job = job.waitFor();

  // Check the table
  StandardTableDefinition table = bigquery.getTable(destinationTable).getDefinition();
  System.out.println("State: " + job.getStatus().getState());
  System.out.printf("Copied %d rows.\n", table.getNumRows());
  // [END bigquery_copy_table_multiple_source]
}
 

/** Example of creating a query job. */
// [TARGET create(JobInfo, JobOption...)]
// [VARIABLE "SELECT field FROM my_dataset_name.my_table_name"]
public Job createJob(String query) {
  // [START ]
  Job job = null;
  JobConfiguration jobConfiguration = QueryJobConfiguration.of(query);
  JobInfo jobInfo = JobInfo.of(jobConfiguration);
  try {
    job = bigquery.create(jobInfo);
  } catch (BigQueryException e) {
    // the job was not created
  }
  // [END ]
  return job;
}
 

@Test
public void testFailedLoad() throws Exception {
    when(table.exists()).thenReturn(Boolean.TRUE);
    when(bq.create(ArgumentMatchers.isA(JobInfo.class))).thenReturn(job);
    when(bq.writer(ArgumentMatchers.isA(WriteChannelConfiguration.class))).thenReturn(tableDataWriteChannel);
    when(tableDataWriteChannel.getJob()).thenReturn(job);
    when(job.waitFor(ArgumentMatchers.isA(RetryOption.class))).thenThrow(BigQueryException.class);
    when(job.getStatus()).thenReturn(jobStatus);
    when(job.getStatistics()).thenReturn(stats);

    when(stats.getCreationTime()).thenReturn(0L);
    when(stats.getStartTime()).thenReturn(1L);
    when(stats.getEndTime()).thenReturn(2L);

    final TestRunner runner = buildNewRunner(getProcessor());
    addRequiredPropertiesToRunner(runner);
    runner.assertValid();

    runner.enqueue("{ \"data\": \"datavalue\" }");

    runner.run();

    runner.assertAllFlowFilesTransferred(PutBigQueryBatch.REL_FAILURE);
}
 

@Test
public void testSuccessfulLoad() throws Exception {
    when(table.exists()).thenReturn(Boolean.TRUE);
    when(bq.create(ArgumentMatchers.isA(JobInfo.class))).thenReturn(job);
    when(bq.writer(ArgumentMatchers.isA(WriteChannelConfiguration.class))).thenReturn(tableDataWriteChannel);
    when(tableDataWriteChannel.getJob()).thenReturn(job);
    when(job.waitFor(ArgumentMatchers.isA(RetryOption.class))).thenReturn(job);
    when(job.getStatus()).thenReturn(jobStatus);
    when(job.getStatistics()).thenReturn(stats);

    when(stats.getCreationTime()).thenReturn(0L);
    when(stats.getStartTime()).thenReturn(1L);
    when(stats.getEndTime()).thenReturn(2L);

    final TestRunner runner = buildNewRunner(getProcessor());
    addRequiredPropertiesToRunner(runner);
    runner.assertValid();

    runner.enqueue("{ \"data\": \"datavalue\" }");

    runner.run();

    runner.assertAllFlowFilesTransferred(PutBigQueryBatch.REL_SUCCESS);
}
 

TableInfo createTableFromQuery()
{
    TableId destinationTable = bigQueryClient.createDestinationTable(table);
    log.debug("destinationTable is %s", destinationTable);
    JobInfo jobInfo = JobInfo.of(
            QueryJobConfiguration
                    .newBuilder(query)
                    .setDestinationTable(destinationTable)
                    .build());
    log.debug("running query %s", jobInfo);
    Job job = waitForJob(bigQueryClient.create(jobInfo));
    log.debug("job has finished. %s", job);
    if (job.getStatus().getError() != null) {
        throw convertToBigQueryException(job.getStatus().getError());
    }
    // add expiration time to the table
    TableInfo createdTable = bigQueryClient.getTable(destinationTable);
    long expirationTime = createdTable.getCreationTime() +
            TimeUnit.HOURS.toMillis(config.viewExpirationTimeInHours);
    Table updatedTable = bigQueryClient.update(createdTable.toBuilder()
            .setExpirationTime(expirationTime)
            .build());
    return updatedTable;
}
 

@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(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 running a batch query. */
public void runBatchQuery() throws TimeoutException, InterruptedException {
  // [START bigquery_query_batch]
  // BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
  String query = "SELECT corpus FROM `bigquery-public-data.samples.shakespeare` GROUP BY corpus;";
  QueryJobConfiguration queryConfig =
      QueryJobConfiguration.newBuilder(query)
          // Run at batch priority, which won't count toward concurrent rate
          // limit.
          .setPriority(QueryJobConfiguration.Priority.BATCH)
          .build();

  // Location must match that of the dataset(s) referenced in the query.
  JobId jobId = JobId.newBuilder().setRandomJob().setLocation("US").build();
  String jobIdString = jobId.getJob();

  // API request - starts the query.
  bigquery.create(JobInfo.newBuilder(queryConfig).setJobId(jobId).build());

  // Check on the progress by getting the job's updated state. Once the state
  // is `DONE`, the results are ready.
  Job queryJob =
      bigquery.getJob(JobId.newBuilder().setJob(jobIdString).setLocation("US").build());
  System.out.printf(
      "Job %s in location %s currently in state: %s%n",
      queryJob.getJobId().getJob(),
      queryJob.getJobId().getLocation(),
      queryJob.getStatus().getState().toString());
  // [END bigquery_query_batch]
}
 

private void generateTableWithDdl(String datasetId, String tableId) throws InterruptedException {
  String sql =
      String.format(
          "CREATE TABLE %s.%s "
              + "AS "
              + "SELECT "
              + "2000 + CAST(18 * RAND() as INT64) AS year, "
              + "IF(RAND() > 0.5,\"foo\",\"bar\") AS token "
              + "FROM "
              + "UNNEST(GENERATE_ARRAY(0,5,1)) AS r",
          datasetId, tableId);
  Job job = bigquery.create(JobInfo.of(QueryJobConfiguration.newBuilder(sql).build()));
  job.waitFor();
}
 

@Test
public void testCancel() {
  JobConfiguration jobConfig =
      QueryJobConfiguration.newBuilder(QUERY).setUseLegacySql(false).build();
  JobInfo jobInfo = JobInfo.newBuilder(jobConfig).build();
  Job job = bigquery.create(jobInfo);
  JobSnippets jobSnippets = new JobSnippets(job);
  boolean result = jobSnippets.cancel();
  assertTrue(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 testIsDone() 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);
  jobSnippets.isDone();
  assertTrue(job.isDone());
}
 

@Test
public void testWaitForWithOptions() 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);
  boolean result = jobSnippets.waitForWithOptions();
  assertTrue(result);
}
 

/** Example of undeleting a table. */
public void undeleteTable(String datasetId) throws InterruptedException {
  generateTableWithDdl(datasetId, "oops_undelete_me");

  // [START bigquery_undelete_table]
  // String datasetId = "my_dataset";
  String tableId = "oops_undelete_me";

  // Record the current time.  We'll use this as the snapshot time
  // for recovering the table.
  long snapTime = Instant.now().toEpochMilli();

  // "Accidentally" delete the table.
  bigquery.delete(TableId.of(datasetId, tableId));

  // Construct the restore-from tableID using a snapshot decorator.
  String snapshotTableId = String.format("%s@%d", tableId, snapTime);
  // Choose a new table ID for the recovered table data.
  String recoverTableId = String.format("%s_recovered", tableId);

  // Construct and run a copy job.
  CopyJobConfiguration configuration =
      CopyJobConfiguration.newBuilder(
              TableId.of(datasetId, recoverTableId), TableId.of(datasetId, snapshotTableId))
          .build();
  Job job = bigquery.create(JobInfo.of(configuration));
  job = job.waitFor();

  // Check the table
  StandardTableDefinition table =
      bigquery.getTable(TableId.of(datasetId, recoverTableId)).getDefinition();
  System.out.println("State: " + job.getStatus().getState());
  System.out.printf("Recovered %d rows.\n", table.getNumRows());
  // [END bigquery_undelete_table]
}
 

/** Example of loading a newline-delimited-json file with textual fields from GCS to a table. */
// [TARGET create(JobInfo, JobOption...)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public Long writeRemoteFileToTable(String datasetName, String tableName)
    throws InterruptedException {
  // [START bigquery_load_table_gcs_json]
  String sourceUri = "gs://cloud-samples-data/bigquery/us-states/us-states.json";
  TableId tableId = TableId.of(datasetName, tableName);
  // Table field definition
  Field[] fields =
      new Field[] {
        Field.of("name", LegacySQLTypeName.STRING),
        Field.of("post_abbr", LegacySQLTypeName.STRING)
      };
  // Table schema definition
  Schema schema = Schema.of(fields);
  LoadJobConfiguration configuration =
      LoadJobConfiguration.builder(tableId, sourceUri)
          .setFormatOptions(FormatOptions.json())
          .setCreateDisposition(CreateDisposition.CREATE_IF_NEEDED)
          .setSchema(schema)
          .build();
  // Load the table
  Job loadJob = bigquery.create(JobInfo.of(configuration));
  loadJob = loadJob.waitFor();
  // Check the table
  System.out.println("State: " + loadJob.getStatus().getState());
  return ((StandardTableDefinition) bigquery.getTable(tableId).getDefinition()).getNumRows();
  // [END bigquery_load_table_gcs_json]
}
 

@Test
public void testWaitFor() 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);
  boolean result = jobSnippets.waitFor();
  assertTrue(result);
}
 

@Test(expected = BigQueryException.class)
public void shouldFailWhenJobTerminatesExceptionally() throws InterruptedException {
  when(bigQuery.getDataset(DATASET_ID)).thenReturn(mock(Dataset.class));

  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  doThrow(new BigQueryException(mock(IOException.class))).when(job)
      .waitFor(any(RetryOption.class));

  BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID).provide(null).publish();
}
 

@Override
JobInfo parse(String... args) throws Exception {
  if (args.length >= 4) {
    String dataset = args[0];
    String table = args[1];
    String format = args[2];
    TableId tableId = TableId.of(dataset, table);
    LoadJobConfiguration configuration =
        LoadJobConfiguration.of(
            tableId, Arrays.asList(args).subList(3, args.length), FormatOptions.of(format));
    return JobInfo.of(configuration);
  }
  throw new IllegalArgumentException("Missing required arguments.");
}
 

@Override
JobInfo parse(String... args) throws Exception {
  if (args.length >= 4) {
    String dataset = args[0];
    String table = args[1];
    String format = args[2];
    TableId tableId = TableId.of(dataset, table);
    ExtractJobConfiguration configuration =
        ExtractJobConfiguration.of(
            tableId, Arrays.asList(args).subList(3, args.length), format);
    return JobInfo.of(configuration);
  }
  throw new IllegalArgumentException("Missing required arguments.");
}
 

@Override
JobInfo parse(String... args) throws Exception {
  String message;
  if (args.length == 4) {
    TableId sourceTableId = TableId.of(args[0], args[1]);
    TableId destinationTableId = TableId.of(args[2], args[3]);
    return JobInfo.of(CopyJobConfiguration.of(destinationTableId, sourceTableId));
  } else if (args.length < 3) {
    message = "Missing required source or destination table.";
  } else {
    message = "Too many arguments.";
  }
  throw new IllegalArgumentException(message);
}
 

@Test
public void testExists() 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);
  boolean result = jobSnippets.exists();
  assertTrue(result);
}
 

@Test(expected = RuntimeException.class)
public void shouldFailWhenJobDisappears() 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(null);

  BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID).provide(null).publish();
}