com.google.api.services.bigquery.Bigquery Java Examples

@Experimental(Kind.SCHEMAS)
public TableDataInsertAllResponse insertRows(Schema rowSchema, Row... rows) throws IOException {
  List<Rows> bqRows =
      Arrays.stream(rows)
          .map(row -> new Rows().setJson(BigQueryUtils.toTableRow(row)))
          .collect(ImmutableList.toImmutableList());
  Bigquery bq = newBigQueryClient(pipelineOptions);

  return bq.tabledata()
      .insertAll(
          pipelineOptions.getProject(),
          pipelineOptions.getTargetDataset(),
          table.getTableReference().getTableId(),
          new TableDataInsertAllRequest().setRows(bqRows))
      .execute();
}
 

@Override
public void cancel(InterpreterContext context) {
  logger.info("Trying to Cancel current query statement.");

  if (service != null && jobId != null && projectId != null) {
    try {
      Bigquery.Jobs.Cancel request = service.jobs().cancel(projectId, jobId);
      JobCancelResponse response = request.execute();
      jobId = null;
      logger.info("Query Execution cancelled");
    } catch (IOException ex) {
      logger.error("Could not cancel the SQL execution");
    }
  } else {
    logger.info("Query Execution was already cancelled");
  }
}
 

@After
public void tearDown() throws IOException {
  // Delete the test dataset along with all tables inside it.
  // TODO(user): Move this into library shared by BigQueryOutputCommitter.
  Bigquery.Datasets datasets = bigqueryInstance.datasets();
  logger.atInfo().log(
      "Deleting temporary test dataset '%s' for project '%s'", testDataset, projectIdValue);
  datasets.delete(projectIdValue, testDataset).setDeleteContents(true).execute();

  // Recursively delete the testBucket.
  setConfigForGcsFromBigquerySettings();
  Path toDelete = new Path(String.format("gs://%s", testBucket));
  FileSystem fs = toDelete.getFileSystem(config);
  if ("gs".equals(fs.getScheme())) {
    bucketHelper.cleanup(
        GoogleCloudStorageFileSystemIntegrationHelper.createGcsFs(projectIdValue).getGcs());
  } else {
    logger.atInfo().log("Deleting temporary test bucket '%s'", toDelete);
    fs.delete(toDelete, true);
  }
}
 

private static void updateTable(Bigquery bigquery, Table table) throws IOException {
  TableReference ref = table.getTableReference();
  try {
    bigquery
        .tables()
        .update(ref.getProjectId(), ref.getDatasetId(), ref.getTableId(), table)
        .execute();
  } catch (GoogleJsonResponseException e) {
    if (e.getDetails() != null && e.getDetails().getCode() == 404) {
      bigquery.tables().insert(ref.getProjectId(), ref.getDatasetId(), table).execute();
    } else {
      logger.atWarning().withCause(e).log(
          "UpdateSnapshotViewAction failed, caught exception %s", e.getDetails());
    }
  }
}
 

/** Ensures the table exists in Bigquery. */
private void ensureTable(Bigquery bigquery, TableReference table, List<TableFieldSchema> schema)
    throws IOException {
  try {
    bigquery.tables().insert(table.getProjectId(), table.getDatasetId(), new Table()
        .setSchema(new TableSchema().setFields(schema))
        .setTableReference(table))
        .execute();
    logger.atInfo().log(
        "Created BigQuery table %s:%s.%s",
        table.getProjectId(), table.getDatasetId(), table.getTableId());
  } catch (IOException e) {
    // Swallow errors about a table that exists, and throw any other ones.
    if (!BigqueryJobFailureException.create(e).getReason().equals("duplicate")) {
      throw e;
    }
  }
}
 

/**
 * Ensures the dataset exists by trying to create it. Note that it's not appreciably cheaper
 * to check for dataset existence than it is to try to create it and check for exceptions.
 */
// Note that these are not static so they can be mocked for testing.
private void ensureDataset(Bigquery bigquery, String projectId, String datasetId)
    throws IOException {
  try {
    bigquery.datasets()
        .insert(projectId,
            new Dataset().setDatasetReference(
                new DatasetReference()
                    .setProjectId(projectId)
                    .setDatasetId(datasetId)))
        .execute();
  } catch (IOException e) {
    // Swallow errors about a duplicate dataset, and throw any other ones.
    if (!BigqueryJobFailureException.create(e).getReason().equals("duplicate")) {
      throw e;
    }
  }
}
 

/**
 * Returns a new connection to Bigquery, first ensuring that the given dataset and table exist in
 * project with the given id, creating them if required.
 */
public Bigquery ensureDataSetAndTableExist(String projectId, String datasetId, String tableId)
    throws IOException {
  ensureDataSetExists(projectId, datasetId);
  checkArgument(bigquerySchemas.containsKey(tableId), "Unknown table ID: %s", tableId);

  if (!knownExistingTables.contains(tableId)) {
    ensureTable(
        bigquery,
        new TableReference()
            .setDatasetId(datasetId)
            .setProjectId(projectId)
            .setTableId(tableId),
          bigquerySchemas.get(tableId));
    knownExistingTables.add(tableId);
  }

  return bigquery;
}
 

static List<DatasetList.Datasets> listDatasets(Bigquery bq, String projectId, Predicate<DatasetList.Datasets> needle)
        throws IOException, RetryExecutor.RetryGiveupException
{
    List<DatasetList.Datasets> datasets = new ArrayList<>();
    Bigquery.Datasets.List req = retryExecutor.run(() -> bq.datasets().list(projectId));
    DatasetList datasetList;
    do {
        datasetList = req.execute();
        if (datasetList.getDatasets() != null) {
            datasetList.getDatasets().stream()
                    .filter(needle)
                    .forEach(datasets::add);
        }
        req.setPageToken(datasetList.getNextPageToken());
    }
    while (null != datasetList.getNextPageToken());
    return datasets;
}
 

static List<TableList.Tables> listTables(Bigquery bq, String projectId, String datasetId, Predicate<TableList.Tables> needle)
        throws IOException
{
    List<TableList.Tables> tables = new ArrayList<>();
    Bigquery.Tables.List req = bq.tables().list(projectId, datasetId);
    TableList tableList;
    do {
        tableList = req.execute();
        if (tableList.getTables() != null) {
            tableList.getTables().stream().filter(needle).forEach(tables::add);
        }
        req.setPageToken(tableList.getNextPageToken());
    }
    while (null != tableList.getNextPageToken());
    return tables;
}
 

/** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */
private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
  RetryHttpRequestInitializer httpRequestInitializer =
      new RetryHttpRequestInitializer(ImmutableList.of(404));
  httpRequestInitializer.setCustomErrors(createBigQueryClientCustomErrors());
  httpRequestInitializer.setWriteTimeout(options.getHTTPWriteTimeout());
  return new Bigquery.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              httpRequestInitializer))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

@Before
public void setUp() {
  MockitoAnnotations.initMocks(this);

  // Set up the MockHttpRequest for future inspection
  request =
      new MockLowLevelHttpRequest() {
        @Override
        public LowLevelHttpResponse execute() throws IOException {
          return response;
        }
      };

  // A mock transport that lets us mock the API responses.
  MockHttpTransport transport =
      new MockHttpTransport.Builder().setLowLevelHttpRequest(request).build();

  // A sample BigQuery API client that uses default JsonFactory and RetryHttpInitializer.
  bigquery =
      new Bigquery.Builder(
              transport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())
          .build();
}
 

static Table createTable(Bigquery bq, String projectId, String datasetId, String tableId)
        throws IOException
{
    Table table = new Table()
            .setTableReference(new TableReference()
                    .setProjectId(projectId)
                    .setDatasetId(datasetId)
                    .setTableId(tableId));
    Table created = createTable(bq, projectId, datasetId, table);
    assertThat(tableExists(bq, projectId, datasetId, tableId), is(true));
    return created;
}
 

static boolean datasetExists(Bigquery bq, String projectId, String datasetId)
        throws IOException
{
    try {
        bq.datasets().get(projectId, datasetId).execute();
        return true;
    }
    catch (GoogleJsonResponseException e) {
        if (e.getStatusCode() == HttpStatusCodes.STATUS_CODE_NOT_FOUND) {
            return false;
        }
        throw e;
    }
}
 

static void cleanupBq(Bigquery bq, String gcpProjectId)
        throws IOException, RetryExecutor.RetryGiveupException
{
    if (bq == null) {
        return;
    }
    List<DatasetList.Datasets> datasets = listDatasets(
            bq, gcpProjectId, ds -> ds.getDatasetReference().getDatasetId().contains(BQ_TAG));
    for (DatasetList.Datasets dataset : datasets) {
        deleteDataset(bq, gcpProjectId, dataset.getDatasetReference().getDatasetId());
    }
}
 

private Bigquery bqClient(GoogleCredential credential)
{
    if (credential.createScopedRequired()) {
        credential = credential.createScoped(BigqueryScopes.all());
    }
    return new Bigquery.Builder(transport, jsonFactory, credential)
            .setApplicationName("digdag-test")
            .build();
}
 

@Override
protected Bigquery client(GoogleCredential credential, HttpTransport transport, JsonFactory jsonFactory)
{
    if (credential.createScopedRequired()) {
        credential = credential.createScoped(BigqueryScopes.all());
    }

    return new Bigquery.Builder(transport, jsonFactory, credential)
            .setApplicationName("Digdag")
            .build();
}
 

public static Job pollJob(final Bigquery.Jobs.Get request, final long interval) 
    throws IOException, InterruptedException {
  Job job = request.execute();
  while (!job.getStatus().getState().equals("DONE")) {
    System.out.println("Job is " 
        + job.getStatus().getState() 
        + " waiting " + interval + " milliseconds...");
    Thread.sleep(interval);
    job = request.execute();
  }
  return job;
}
 

static boolean tableExists(Bigquery bq, String projectId, String datasetId, String tableId)
        throws IOException
{
    try {
        bq.tables().get(projectId, datasetId, tableId).execute();
        return true;
    }
    catch (GoogleJsonResponseException e) {
        if (e.getStatusCode() == HttpStatusCodes.STATUS_CODE_NOT_FOUND) {
            return false;
        }
        throw e;
    }
}
 

static Dataset createDataset(Bigquery bq, String projectId, String datasetId)
        throws IOException, RetryExecutor.RetryGiveupException
{
    Dataset dataset = new Dataset()
            .setDatasetReference(new DatasetReference()
                    .setDatasetId(datasetId));
    Dataset created = createDataset(bq, projectId, dataset);
    assertThat(datasetExists(bq, projectId, datasetId), is(true));
    return created;
}
 

@Before
public void setUp() throws Exception {
    configuration = new GoogleBigQueryConfiguration();

    bigquery = Mockito.mock(Bigquery.class);
    endpoint = Mockito.mock(GoogleBigQueryEndpoint.class);
    tabledata = Mockito.mock(Bigquery.Tabledata.class);
    mockInsertall = Mockito.mock(Bigquery.Tabledata.InsertAll.class);

    Mockito.when(bigquery.tabledata()).thenReturn(tabledata);
    Mockito.when(tabledata.insertAll(Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.any())).thenReturn(mockInsertall);

    TableDataInsertAllResponse mockResponse = new TableDataInsertAllResponse();
    Mockito.when(mockInsertall.execute()).thenReturn(mockResponse);
}
 

/**
 * Constructs a BigQuery from the credential constructed from the environment.
 *
 * @throws IOException on IO Error.
 * @throws GeneralSecurityException on General Security Error.
 */
public Bigquery getBigQuery(Configuration config)
    throws GeneralSecurityException, IOException {
  logger.atInfo().log("Creating BigQuery from default credential.");
  Credential credential = createBigQueryCredential(config);
  // Use the credential to create an authorized BigQuery client
  return getBigQueryFromCredential(config, credential, BQC_ID);
}
 

/** Constructs a BigQuery from a given Credential. */
public Bigquery getBigQueryFromCredential(
    Configuration config, Credential credential, String appName) {
  logger.atInfo().log("Creating BigQuery from given credential.");
  // Use the credential to create an authorized BigQuery client
  if (credential != null) {
    return new Bigquery
        .Builder(HTTP_TRANSPORT, JSON_FACTORY, new RetryHttpInitializer(credential, appName))
        .setApplicationName(appName).build();
  }
  return new Bigquery.Builder(HTTP_TRANSPORT, JSON_FACTORY, /* httpRequestInitializer= */ null)
      .setRootUrl(BQ_ROOT_URL.get(config, config::get))
      .setApplicationName(appName)
      .build();
}
 

/**
 * Cleans up relevant temporary resources associated with a job which used the
 * GsonBigQueryInputFormat; this should be called explicitly after the completion of the entire
 * job. Possibly cleans up intermediate export tables if configured to use one due to
 * specifying a BigQuery "query" for the input. Cleans up the GCS directoriy where BigQuery
 * exported its files for reading.
 */
public static void cleanupJob(Configuration configuration, JobID jobId) throws IOException {
  String exportPathRoot = BigQueryConfiguration.getTemporaryPathRoot(configuration, jobId);
  configuration.set(TEMP_GCS_PATH.getKey(), exportPathRoot);
  Bigquery bigquery;
  try {
    bigquery = new BigQueryFactory().getBigQuery(configuration);
  } catch (GeneralSecurityException gse) {
    throw new IOException("Failed to create Bigquery client", gse);
  }
  cleanupJob(new BigQueryHelper(bigquery), configuration);
}
 

@Test
public void getBigQuery_localDev() throws Exception {
  BigQueryFactory factory = new BigQueryFactory();
  Configuration configuration = CredentialConfigurationUtil.getTestConfiguration();

  Bigquery bigquery = factory.getBigQuery(configuration);

  assertThat(bigquery).isNotNull();
  assertThat(bigquery.getRootUrl()).isEqualTo("https://bigquery.googleapis.com/");
}
 

@Test
public void getBigQuery_customEndpoint() throws Exception {
  BigQueryFactory factory = new BigQueryFactory();
  Configuration configuration = CredentialConfigurationUtil.getTestConfiguration();
  configuration.set(BQ_ROOT_URL.getKey(), "https://unit-test-bigquery.googleapis.com/");

  Bigquery bigquery = factory.getBigQuery(configuration);

  assertThat(bigquery).isNotNull();
  assertThat(bigquery.getRootUrl()).isEqualTo("https://unit-test-bigquery.googleapis.com/");
}
 

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

/** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */
private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
  return new Bigquery.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              new RetryHttpRequestInitializer(ImmutableList.of(404))))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

private String uploadBackup(String backupId, String backupFolderUrl, Iterable<String> kinds)
    throws IOException {
  Bigquery bigquery = checkedBigquery.ensureDataSetExists(projectId, BACKUP_DATASET);
  String loadMessage =
      String.format("Loading Datastore backup %s from %s...", backupId, backupFolderUrl);
  logger.atInfo().log(loadMessage);

  String sanitizedBackupId = sanitizeForBigquery(backupId);
  StringBuilder builder = new StringBuilder(loadMessage + "\n");
  builder.append("Load jobs:\n");

  for (String kindName : kinds) {
    String jobId = String.format("load-backup-%s-%s", sanitizedBackupId, kindName);
    JobReference jobRef = new JobReference().setProjectId(projectId).setJobId(jobId);
    String sourceUri = getBackupInfoFileForKind(backupFolderUrl, kindName);
    String tableId = String.format("%s_%s", sanitizedBackupId, kindName);

    // Launch the load job.
    Job job = makeLoadJob(jobRef, sourceUri, tableId);
    bigquery.jobs().insert(projectId, job).execute();

    // Enqueue a task to check on the load job's completion, and if it succeeds, to update a
    // well-known view in BigQuery to point at the newly loaded backup table for this kind.
    bigqueryPollEnqueuer.enqueuePollTask(
        jobRef,
        createViewUpdateTask(BACKUP_DATASET, tableId, kindName, LATEST_BACKUP_VIEW_NAME),
        getQueue(UpdateSnapshotViewAction.QUEUE));

    builder.append(String.format(" - %s:%s\n", projectId, jobId));
    logger.atInfo().log("Submitted load job %s:%s", projectId, jobId);
  }
  return builder.toString();
}
 

/** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */
private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) {
  return new Bigquery.Builder(
          Transport.getTransport(),
          Transport.getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              new RetryHttpRequestInitializer(ImmutableList.of(404))))
      .setApplicationName(options.getAppName())
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

private TableSchema getSchema(Bigquery bq) {
  try {
    return bq.tables()
        .get(
            pipelineOptions.getProject(),
            pipelineOptions.getTargetDataset(),
            table.getTableReference().getTableId())
        .execute()
        .getSchema();
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
}