com.google.cloud.bigquery.BigQuery Java Examples

public static void main(String... args) throws InterruptedException, TimeoutException {
  BigQuery bigquery = BigQueryOptions.getDefaultInstance().getService();
  TableId tableId = TableId.of("dataset", "table");
  Table table = bigquery.getTable(tableId);
  if (table == null) {
    System.out.println("Creating table " + tableId);
    Field integerField = Field.of("fieldName", LegacySQLTypeName.INTEGER);
    Schema schema = Schema.of(integerField);
    table = bigquery.create(TableInfo.of(tableId, StandardTableDefinition.of(schema)));
  }
  System.out.println("Loading data into table " + tableId);
  Job loadJob = table.load(FormatOptions.csv(), "gs://bucket/path");
  loadJob = loadJob.waitFor();
  if (loadJob.getStatus().getError() != null) {
    System.out.println("Job completed with errors");
  } else {
    System.out.println("Job succeeded");
  }
}
 

private TableResult queryWithLarge(BigQuery bigquery, QueryJobConfiguration queryRequest, String projectId,
                                   BigQuery.JobOption... options) {
    String tempDataset = genTempName("dataset");
    String tempTable = genTempName("table");
    bigquery.create(DatasetInfo.of(tempDataset));
    TableId tableId = TableId.of(projectId, tempDataset, tempTable);
    QueryJobConfiguration jobConfiguration = QueryJobConfiguration
            .newBuilder(queryRequest.getQuery())
            .setAllowLargeResults(true)
            .setUseLegacySql(queryRequest.useLegacySql())
            .setDestinationTable(tableId)
            .build();
    try {
        return query(bigquery, jobConfiguration, projectId, options);
    } finally {
        bigquery.delete(tableId);
    }
}
 

@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 testSettingBigQueryOptions() {
	this.contextRunner.run((context) -> {
		BigQueryOptions bigQueryOptions = context.getBean(BigQuery.class).getOptions();
		assertThat(bigQueryOptions.getProjectId()).isEqualTo("test-project");
		assertThat(bigQueryOptions.getCredentials()).isEqualTo(MOCK_CREDENTIALS);

		BigQueryTemplate bigQueryTemplate = context.getBean(BigQueryTemplate.class);
		assertThat(bigQueryTemplate.getDatasetName()).isEqualTo("test-dataset");
	});
}
 

/**
 * Creates the {@link BigQuery} template.
 *
 * @param bigQuery the underlying client object used to interface with BigQuery
 * @param datasetName the name of the dataset in which all operations will take place
 * @param taskScheduler the {@link TaskScheduler} used to poll for the status of
 *     long-running BigQuery operations
 */
public BigQueryTemplate(BigQuery bigQuery, String datasetName, TaskScheduler taskScheduler) {
	Assert.notNull(bigQuery, "BigQuery client object must not be null.");
	Assert.notNull(datasetName, "Dataset name must not be null");
	Assert.notNull(taskScheduler, "TaskScheduler must not be null");

	this.bigQuery = bigQuery;
	this.datasetName = datasetName;
	this.taskScheduler = taskScheduler;
}
 

@Bean
@ConditionalOnMissingBean
public BigQuery bigQuery() throws IOException {
	BigQueryOptions bigQueryOptions = BigQueryOptions.newBuilder()
			.setProjectId(this.projectId)
			.setCredentials(this.credentialsProvider.getCredentials())
			.setHeaderProvider(new UserAgentHeaderProvider(GcpBigQueryAutoConfiguration.class))
			.build();
	return bigQueryOptions.getService();
}
 

public ProtoUpdateListener(ProtoMappingConfig protoMappingConfig, StencilConfig stencilConfig, BQConfig bqConfig, Converter protoMappingConverter, Parser protoMappingParser, BigQuery bqInstance) throws IOException {
    super(stencilConfig.getProtoSchema());
    this.proto = stencilConfig.getProtoSchema();
    this.protoMappingConfig = protoMappingConfig;
    this.stencilConfig = stencilConfig;
    this.protoMappingConverter = protoMappingConverter;
    this.protoMappingParser = protoMappingParser;
    this.protoFieldFactory = new ProtoFieldFactory();
    this.bqClient = new BQClient(bqInstance, bqConfig);
    this.createStencilClient();
    this.setProtoParser(getProtoMapping());
}
 

public void getFields() {
  try {

    BeamBQInputMeta meta = new BeamBQInputMeta();
    getInfo(meta);

    BigQuery bigQuery = BigQueryOptions.getDefaultInstance().getService();

    if ( StringUtils.isNotEmpty( meta.getDatasetId() ) &&
         StringUtils.isNotEmpty( meta.getTableId() )) {

      Table table = bigQuery.getTable(
        transMeta.environmentSubstitute( meta.getDatasetId()),
        transMeta.environmentSubstitute( meta.getTableId() )
      );

      TableDefinition definition = table.getDefinition();
      Schema schema = definition.getSchema();
      FieldList fieldList = schema.getFields();

      RowMetaInterface rowMeta = new RowMeta();
      for ( int i = 0; i< fieldList.size(); i++) {
        Field field = fieldList.get( i );

        String name = field.getName();
        String type = field.getType().name();

        int kettleType = BQSchemaAndRecordToKettleFn.AvroType.valueOf( type ).getKettleType();
        rowMeta.addValueMeta( ValueMetaFactory.createValueMeta( name, kettleType ) );
      }

      BaseStepDialog.getFieldsFromPrevious( rowMeta, wFields, 1, new int[] { 1 }, new int[] { 3 }, -1, -1, true, null );
    }

  } catch ( Exception e ) {
    new ErrorDialog( shell, "Error", "Error getting BQ fields", e );
  }
}
 

@Override
void run(BigQuery bigquery, QueryJobConfiguration queryConfig) throws Exception {
  System.out.println("Running query");
  for (FieldValueList row : bigquery.query(queryConfig).iterateAll()) {
    System.out.println(row);
  }
}
 

@Override
public void run(BigQuery bigquery, TableId tableId) {
  if (bigquery.delete(tableId)) {
    System.out.println("Table " + tableId + " was deleted");
  } else {
    System.out.println("Table " + tableId + " not found");
  }
}
 

@Override
void run(BigQuery bigquery, Tuple<WriteChannelConfiguration, String> configuration)
    throws Exception {
  System.out.println("Running insert");
  try (FileChannel fileChannel = FileChannel.open(Paths.get(configuration.y()))) {
    WriteChannel writeChannel = bigquery.writer(configuration.x());
    long position = 0;
    long written = fileChannel.transferTo(position, CHUNK_SIZE, writeChannel);
    while (written > 0) {
      position += written;
      written = fileChannel.transferTo(position, CHUNK_SIZE, writeChannel);
    }
    writeChannel.close();
  }
}
 

@Override
public void run(BigQuery bigquery, JobId jobId) {
  if (bigquery.cancel(jobId)) {
    System.out.println("Requested cancel for job " + jobId);
  } else {
    System.out.println("Job " + jobId + " not found");
  }
}
 

private BigQuery getBigQueryInstance() {
    final TransportOptions transportOptions = BigQueryOptions.getDefaultHttpTransportOptions().toBuilder()
            .setConnectTimeout(Integer.parseInt(bqConfig.getBqClientConnectTimeout()))
            .setReadTimeout(Integer.parseInt(bqConfig.getBqClientReadTimeout()))
            .build();
    return BigQueryOptions.newBuilder()
            .setTransportOptions(transportOptions)
            .setCredentials(getGoogleCredentials())
            .setProjectId(bqConfig.getGCPProject())
            .build().getService();
}
 

private Sink createBigQuerySink() {
    BigQuery bq = getBigQueryInstance();
    BQResponseParser responseParser = new BQResponseParser();
    BQErrorHandler bqErrorHandler = createOOBErrorHandler();
    BQRow recordInserter = new BQRowWithInsertId();
    if (!bqConfig.isBQRowInsertIdEnabled()) {
        recordInserter = new BQRowWithoutId();
    }
    Sink bqSink = new BqSink(bq, TableId.of(bqConfig.getDataset(), bqConfig.getTable()), responseParser, bqErrorHandler, recordInserter);
    return new RetrySink(bqSink, new ExponentialBackOffProvider(backOffConfig.getExponentialBackoffInitialTimeInMs(), backOffConfig.getExponentialBackoffMaximumTimeInMs(), backOffConfig.getExponentialBackoffRate(), new BackOff()), appConfig.getMaxPushAttempts());
}
 

@Bean
public BigQuery getInstance() throws IOException {
	// projectId needs to be set explicitly even if it's there in the json key!!
	BigQuery bigQuery = BigQueryOptions.newBuilder().setProjectId(porjectIdProvider.getProjectId())
			.setCredentials(credentialsProvider.getCredentials()).build().getService();

	// Use the client.
	LOGGER.info("Datasets:");
	for (Dataset dataset : bigQuery.listDatasets().iterateAll()) {
		LOGGER.info(dataset.getDatasetId().getDataset());
	}
	return bigQuery;

}
 

@Override
public AbstractBigQueryProcessor getProcessor() {
    return new PutBigQueryBatch() {
        @Override
        protected BigQuery getCloudService() {
            return bq;
        }
    };
}
 

@AfterClass
public static void tearDown() {
  BigQueryOptions options = BigQueryOptions.newBuilder().build();
  BigQuery client = options.getService();
  TableId tableId = TableId.of(options.getProjectId(), testBigQueryDataset, testBigQueryTable);
  client.delete(tableId);
}
 

@Override
public void run(BigQuery bigquery, DatasetId datasetId) {
  if (bigquery.delete(datasetId)) {
    System.out.println("Dataset " + datasetId + " was deleted");
  } else {
    System.out.println("Dataset " + datasetId + " not found");
  }
}
 

/**
 * Constructs a new {@link BigQueryDelegate} for a specific project and set of credentials.
 * #legacySql and #standardSql annotations in the query are not supported, instead the
 * boolean flag useLegacySql should be set to determine which syntax is to be used.
 *
 * @param bigquery BigQuery instance
 * @param useLegacySql When false (default), uses Standard SQL syntax and data types.
 * @param clock Clock implementation for checking whether query has timed out. Normally only supplied in tests.
 */
BigQueryDelegate(BigQuery bigquery, boolean useLegacySql, Clock clock) {
  this.bigquery = bigquery;
  this.clock = clock;

  if (useLegacySql) {
    asRecordFieldTypeFunction = field -> LEGACY_SQL_TYPES.get(field.getType());
    dateTimeFormatter = field -> LegacySQLTypeDateFormatter.get(field.getType());
  } else {
    asRecordFieldTypeFunction = field -> STANDARD_SQL_TYPES.get(field.getType().getStandardType());
    dateTimeFormatter = field -> StandardSQLTypeDateFormatter.get(field.getType().getStandardType());
  }
}
 

public static BigQuery getBigquery(Credentials credentials, String projectId) {
  BigQueryOptions options = BigQueryOptions.newBuilder()
      .setCredentials(credentials)
      .setProjectId(projectId)
      .build();
  return options.getService();
}
 

@Before
public void setUp() throws Exception {
  mockBigquery = mock(BigQuery.class);
  jobId = JobId.of("test-project", "datacollector");
  mockResult = mock(TableResult.class);

  mockStatic(ServiceAccountCredentials.class);
  mockStatic(GoogleCredentials.class);
  ServiceAccountCredentials serviceAccountCredentials = mock(ServiceAccountCredentials.class);
  when(ServiceAccountCredentials.fromStream(any())).thenReturn(serviceAccountCredentials);
}
 

@Override
public Set<String> listDatasets() throws IOException {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    Page<Dataset> datasets = bigquery.listDatasets(properties.getDatastoreProperties().projectName.getValue(),
            BigQuery.DatasetListOption.pageSize(100));
    Set<String> datasetsName = new HashSet<>();
    Iterator<Dataset> datasetIterator = datasets.iterateAll().iterator();
    while (datasetIterator.hasNext()) {
        datasetsName.add(datasetIterator.next().getDatasetId().getDataset());
    }
    return datasetsName;
}
 

@Override
public Set<String> listTables() throws IOException {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    DatasetId datasetId = DatasetId.of(properties.getDatastoreProperties().projectName.getValue(),
            properties.bqDataset.getValue());
    Page<Table> tables = bigquery.listTables(datasetId, BigQuery.TableListOption.pageSize(100));
    Set<String> tablesName = new HashSet<>();
    Iterator<Table> tableIterator = tables.iterateAll().iterator();
    while (tableIterator.hasNext()) {
        tablesName.add(tableIterator.next().getTableId().getTable());
    }
    return tablesName;
}
 

/**
 * Get the schema by table name or query. This method also needed for read and write, because we can not get schema
 * from the ongoing data. BigQueryIO.Read return TableRow, which do not include schema in itself. So use BigQuery
 * client to get it before read and write.
 * 
 * @return
 */
@Override
public Schema getSchema() {
    BigQuery bigquery = BigQueryConnection.createClient(properties.getDatastoreProperties());
    com.google.cloud.bigquery.Schema bqRowSchema = null;
    switch (properties.sourceType.getValue()) {
    case TABLE_NAME: {
        TableId tableId = TableId.of(properties.getDatastoreProperties().projectName.getValue(),
                properties.bqDataset.getValue(), properties.tableName.getValue());
        Table table = bigquery.getTable(tableId);
        if (table == null) {
            ComponentException.build(CommonErrorCodes.UNEXPECTED_EXCEPTION).setAndThrow(
                    "Table not found:" + tableId.toString());
        }
        bqRowSchema = table.getDefinition().getSchema();
        break;
    }
    case QUERY: {
        QueryJobConfiguration queryRequest = QueryJobConfiguration
                .newBuilder(properties.query.getValue())
                .setUseLegacySql(properties.useLegacySql.getValue())
                .build();
        TableResult queryResponse =
                query(bigquery, queryRequest, properties.getDatastoreProperties().projectName.getValue());
        bqRowSchema = queryResponse.getSchema();
        break;
    }
    default:
        throw new RuntimeException("To be implemented: " + properties.sourceType.getValue());
    }
    return BigQueryAvroRegistry.get().inferSchema(bqRowSchema);
}
 

private TableResult query(BigQuery bigquery, QueryJobConfiguration queryRequest, String projectId,
                          BigQuery.JobOption... options) {
    TableResult queryResponse = null;
    try {
        queryResponse = bigquery.query(queryRequest, options);
    } catch (BigQueryException exception) {
        if ("responseTooLarge".equals(exception.getReason())) {
            return queryWithLarge(bigquery, queryRequest, projectId, options);
        }
    } catch (final InterruptedException e) {
        Thread.currentThread().interrupt();
    }
    return loopQueryResponse(queryResponse);
}
 

/**
 * Check connection and authentication
 *
 * @param container
 * @return
 */
@Override
public Iterable<ValidationResult> doHealthChecks(RuntimeContainer container) {
    BigQuery bigquery = BigQueryConnection.createClient(properties);
    bigquery.listDatasets(BigQuery.DatasetListOption.pageSize(1));
    return Arrays.asList(ValidationResult.OK);
}
 

public static BigQuery createClient(BigQueryDatastoreProperties datastore) {
    if (StringUtils.isEmpty(datastore.serviceAccountFile.getValue())) {
        return BigQueryOptions.getDefaultInstance().getService();
    } else {
        return BigQueryOptions.newBuilder().setProjectId(datastore.projectName.getValue())
                .setCredentials(createCredentials(datastore)).build().getService();
    }
}
 

@BeforeClass
public static void initDatasetAndTable() throws IOException {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    for (String dataset : datasets) {
        DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, dataset);
        bigquery.create(DatasetInfo.of(datasetId));
    }

    for (String table : tables) {
        TableDefinition tableDefinition =
                StandardTableDefinition.of(Schema.of(Field.of("test", LegacySQLTypeName.STRING)));
        TableId tableId = TableId.of(BigQueryTestConstants.PROJECT, datasets.get(0), table);
        bigquery.create(TableInfo.of(tableId, tableDefinition));
    }
}
 

@AfterClass
public static void cleanDatasetAndTable() {
    BigQuery bigquery = BigQueryConnection.createClient(createDatastore());
    for (String dataset : datasets) {
        DatasetId datasetId = DatasetId.of(BigQueryTestConstants.PROJECT, dataset);
        bigquery.delete(datasetId, BigQuery.DatasetDeleteOption.deleteContents());
    }
}
 

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