com.google.cloud.bigquery.BigQueryOptions Java Examples

@Test
public void testBiqQueryOptionsConfiguration() throws Exception {
    reset(bq);
    final TestRunner runner = buildNewRunner(getProcessor());

    final AbstractBigQueryProcessor processor = getProcessor();
    final GoogleCredentials mockCredentials = mock(GoogleCredentials.class);

    final BigQueryOptions options = processor.getServiceOptions(runner.getProcessContext(),
            mockCredentials);

    assertEquals("Project IDs should match",
            PROJECT_ID, options.getProjectId());

    assertEquals("Retry counts should match",
            RETRIES.intValue(), options.getRetrySettings().getMaxAttempts());

    assertSame("Credentials should be configured correctly",
            mockCredentials, options.getCredentials());
}
 

@BeforeClass
public static void beforeClass() throws IOException {
    final Map<PropertyDescriptor, String> propertiesMap = new HashMap<>();
    propertiesMap.put(CredentialPropertyDescriptors.SERVICE_ACCOUNT_JSON_FILE, SERVICE_ACCOUNT_JSON);
    Credentials credentials = credentialsProviderFactory.getGoogleCredentials(propertiesMap, new ProxyAwareTransportFactory(null));

    BigQueryOptions bigQueryOptions = BigQueryOptions.newBuilder()
            .setProjectId(PROJECT_ID)
            .setCredentials(credentials)
            .build();

    bigquery = bigQueryOptions.getService();

    DatasetInfo datasetInfo = DatasetInfo.newBuilder(RemoteBigQueryHelper.generateDatasetName()).build();
    dataset = bigquery.create(datasetInfo);
}
 

@Override
protected BigQueryOptions getServiceOptions(ProcessContext context, GoogleCredentials credentials) {
    final String projectId = context.getProperty(PROJECT_ID).evaluateAttributeExpressions().getValue();
    final Integer retryCount = Integer.valueOf(context.getProperty(RETRY_COUNT).getValue());

    final BigQueryOptions.Builder builder = BigQueryOptions.newBuilder();

    if (!StringUtils.isBlank(projectId)) {
        builder.setProjectId(projectId);
    }

    return builder.setCredentials(credentials)
            .setRetrySettings(RetrySettings.newBuilder().setMaxAttempts(retryCount).build())
            .setTransportOptions(getTransportOptions(context))
            .build();
}
 

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

@Before
public void setUp() {
  MockitoAnnotations.initMocks(this);
  bout = new ByteArrayOutputStream();
  PrintStream out = new PrintStream(bout);

  MetricServiceClient metricsClient = MetricServiceClient.create(metricsServiceStub);
  app = new BigQueryRunner(metricsClient, BigQueryOptions.getDefaultInstance().getService(), out);

  when(metricsServiceStub.listMetricDescriptorsPagedCallable()).thenReturn(listCallable);
  when(listCallable.call(any(ListMetricDescriptorsRequest.class))).thenReturn(listResponse);
  when(listResponse.iterateAll()).thenReturn(Collections.EMPTY_LIST);

  when(metricsServiceStub.createMetricDescriptorCallable()).thenReturn(createMetricCallable);
  when(createMetricCallable.call(any(CreateMetricDescriptorRequest.class))).thenReturn(null);

  when(metricsServiceStub.createTimeSeriesCallable()).thenReturn(createTimeSeriesCallable);
  when(createTimeSeriesCallable.call(any(CreateTimeSeriesRequest.class)))
      .thenReturn(Empty.getDefaultInstance());
}
 

@BeforeClass
public static void setup() throws IOException {
  options = IOITHelper.readIOTestPipelineOptions(BigQueryPerfTestOptions.class);
  tempRoot = options.getTempRoot();
  sourceOptions =
      SyntheticOptions.fromJsonString(options.getSourceOptions(), SyntheticSourceOptions.class);
  metricsBigQueryDataset = options.getMetricsBigQueryDataset();
  metricsBigQueryTable = options.getMetricsBigQueryTable();
  testBigQueryDataset = options.getTestBigQueryDataset();
  testBigQueryTable = options.getTestBigQueryTable();
  writeFormat = WriteFormat.valueOf(options.getWriteFormat());
  BigQueryOptions bigQueryOptions = BigQueryOptions.newBuilder().build();
  tableQualifier =
      String.format(
          "%s:%s.%s", bigQueryOptions.getProjectId(), testBigQueryDataset, testBigQueryTable);
  settings =
      InfluxDBSettings.builder()
          .withHost(options.getInfluxHost())
          .withDatabase(options.getInfluxDatabase())
          .withMeasurement(options.getInfluxMeasurement())
          .get();
}
 

@Provides
@Singleton
public BigQueryClient provideBigQueryClient(BigQueryConfig config, UserAgentHeaderProvider userAgentHeaderProvider, BigQueryCredentialsSupplier bigQueryCredentialsSupplier) {
    String billingProjectId = calculateBillingProjectId(config.getParentProjectId(), bigQueryCredentialsSupplier.getCredentials());
    BigQueryOptions.Builder options = BigQueryOptions.newBuilder()
            .setHeaderProvider(userAgentHeaderProvider)
            .setProjectId(billingProjectId)
            .setCredentials(bigQueryCredentialsSupplier.getCredentials());
    return new BigQueryClient(
            options.build().getService(),
            config.getMaterializationProject(),
            config.getMaterializationDataset());
}
 

public static void main(String[] args) {
	CreateStore cs = 
		new CreateStore(args[1], BigQueryOptions.newBuilder().setProjectId(args[0]).build().getService());

	String storeName = args[2];
	
	cs.dropStore(storeName);
	cs.createStore(storeName, 0);
}
 

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

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

@Before
public void setUp() {
  bigquery = BigQueryOptions.getDefaultInstance().getService();
  if (bigquery.getDataset(datasetName) == null) {
    bigquery.create(DatasetInfo.newBuilder(datasetName).build());
  }
  bout = new ByteArrayOutputStream();
  out = new PrintStream(bout);
  System.setOut(out);
}
 

@Before
public void setUp() {
  bigquery = BigQueryOptions.getDefaultInstance().getService();
  if (bigquery.getDataset(datasetName) == null) {
    Dataset dataset = bigquery.create(DatasetInfo.newBuilder(datasetName).build());
  }
  bout = new ByteArrayOutputStream();
  out = new PrintStream(bout);
  System.setOut(out);
}
 

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

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

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

private void updateTableIfRequired(TableId tableId, TableRow row,
    Map<String, LegacySQLTypeName> inputSchema) {
  // Ensure Instance of BigQuery Exists
  if (this.bigquery == null) {
    this.bigquery =
        BigQueryOptions.newBuilder()
            .setProjectId(getProjectId())
            .build()
            .getService();
  }

  // Get BigQuery Table for Given Row
  Table table = getBigQueryTable(tableId);

  // Validate Table Schema
  FieldList tableFields = table.getDefinition().getSchema().getFields();

  Set<String> rowKeys = row.keySet();
  Boolean tableWasUpdated = false;
  List<Field> newFieldList = new ArrayList<Field>();
  for (String rowKey : rowKeys) {
    // Check if rowKey (column from data) is in the BQ Table
    try {
      Field tableField = tableFields.get(rowKey);
    } catch (IllegalArgumentException e) {
      tableWasUpdated = addNewTableField(tableId, row, rowKey, newFieldList, inputSchema);
    }
  }

  if (tableWasUpdated) {
    LOG.info("Updating Table");
    updateBigQueryTable(tableId, table, tableFields, newFieldList);
  }
}
 

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

static String calculateBillingProjectId(Optional<String> configParentProjectId, Credentials credentials) {
    // 1. Get from configuration
    if (configParentProjectId.isPresent()) {
        return configParentProjectId.get();
    }
    // 2. Get from the provided credentials, but only ServiceAccountCredentials contains the project id.
    // All other credentials types (User, AppEngine, GCE, CloudShell, etc.) take it from the environment
    if (credentials instanceof ServiceAccountCredentials) {
        return ((ServiceAccountCredentials) credentials).getProjectId();
    }
    // 3. No configuration was provided, so get the default from the environment
    return BigQueryOptions.getDefaultProjectId();
}
 

@Provides
@Singleton
public BigQueryClient provideBigQueryClient(BigQueryConfig config, HeaderProvider headerProvider, BigQueryCredentialsSupplier bigQueryCredentialsSupplier)
{
    String billingProjectId = calculateBillingProjectId(config.getParentProjectId(), bigQueryCredentialsSupplier.getCredentials());
    BigQueryOptions.Builder options = BigQueryOptions.newBuilder()
            .setHeaderProvider(headerProvider)
            .setProjectId(billingProjectId);
    // set credentials of provided
    bigQueryCredentialsSupplier.getCredentials().ifPresent(options::setCredentials);
    return new BigQueryClient(options.build().getService(), config);
}
 

static String calculateBillingProjectId(Optional<String> configParentProjectId, Optional<Credentials> credentials)
{
    // 1. Get from configuration
    if (configParentProjectId.isPresent()) {
        return configParentProjectId.get();
    }
    // 2. Get from the provided credentials, but only ServiceAccountCredentials contains the project id.
    // All other credentials types (User, AppEngine, GCE, CloudShell, etc.) take it from the environment
    if (credentials.isPresent() && credentials.get() instanceof ServiceAccountCredentials) {
        return ((ServiceAccountCredentials) credentials.get()).getProjectId();
    }
    // 3. No configuration was provided, so get the default from the environment
    return BigQueryOptions.getDefaultProjectId();
}
 

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

}
 

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

@Override
public Schema getSchema(TableId tableId, Map<String, String> attributes) {
  if (tableId == null) {
    // Always throws SchemaNotFoundException
    return getSchema(attributes);
  }
  if (tableSchemaCache == null) {
    // We need to be very careful about settings for the cache here. We have had significant
    // issues in the past due to exceeding limits on BigQuery API requests; see
    // https://bugzilla.mozilla.org/show_bug.cgi?id=1623000
    tableSchemaCache = CacheBuilder.newBuilder().expireAfterWrite(Duration.ofMinutes(10))
        .build();
  }
  if (bqService == null) {
    bqService = BigQueryOptions.newBuilder().setProjectId(tableId.getProject())
        .setRetrySettings(RETRY_SETTINGS).build().getService();
  }
  try {
    return Optional.of(tableSchemaCache.get(tableId, () -> {
      Table table = bqService.getTable(tableId);
      if (table != null) {
        return table.getDefinition().getSchema();
      } else {
        return null;
      }
    })).orElseThrow(() -> SchemaNotFoundException.forName(tableId.toString()));
  } catch (ExecutionException e) {
    throw new UncheckedExecutionException(e.getCause());
  }
}
 

/** Sets all objects needed during mapper execution. */
public void setUp() {
  if (this.bqTableRowCleaner == null) {
    this.bqTableRowCleaner = BigQueryTableRowCleaner.getBigQueryTableRowCleaner();
  }
  if (this.bigquery == null) {
    this.bigquery =
        BigQueryOptions.newBuilder().setProjectId(getProjectId()).build().getService();
  }
  if (this.tableCache == null) {
    this.tableCache = new BigQueryTableCache(this.bigquery);
  }
}
 

/**
 * Return the appropriate table destination instance for the given document type and other
 * attributes.
 */
public TableDestination getTableDestination(Map<String, String> attributes) {
  attributes = new HashMap<>(attributes);

  // We coerce all docType and namespace names to be snake_case and to remove invalid
  // characters; these transformations MUST match with the transformations applied by the
  // jsonschema-transpiler and mozilla-schema-generator when creating table schemas in BigQuery.
  final String namespace = attributes.get(Attribute.DOCUMENT_NAMESPACE);
  final String docType = attributes.get(Attribute.DOCUMENT_TYPE);
  if (namespace != null) {
    attributes.put(Attribute.DOCUMENT_NAMESPACE, getAndCacheNormalizedName(namespace));
  }
  if (docType != null) {
    attributes.put(Attribute.DOCUMENT_TYPE, getAndCacheNormalizedName(docType));
  }

  // Only letters, numbers, and underscores are allowed in BigQuery dataset and table names,
  // but some doc types and namespaces contain '-', so we convert to '_'; we don't pass all
  // values through getAndCacheBqName to avoid expensive regex operations and polluting the
  // cache of transformed field names.
  attributes = Maps.transformValues(attributes, v -> v.replaceAll("-", "_"));

  final String tableSpec = StringSubstitutor.replace(tableSpecTemplate.get(), attributes);

  // Send to error collection if incomplete tableSpec; $ is not a valid char in tableSpecs.
  if (tableSpec.contains("$")) {
    throw new IllegalArgumentException("Element did not contain all the attributes needed to"
        + " fill out variables in the configured BigQuery output template: "
        + tableSpecTemplate.get());
  }

  final TableDestination tableDestination = new TableDestination(tableSpec, null,
      new TimePartitioning().setField(partitioningField.get()),
      new Clustering().setFields(clusteringFields.get()));
  final TableReference ref = BigQueryHelpers.parseTableSpec(tableSpec);
  final DatasetReference datasetRef = new DatasetReference().setProjectId(ref.getProjectId())
      .setDatasetId(ref.getDatasetId());

  if (bqService == null) {
    bqService = BigQueryOptions.newBuilder().setProjectId(ref.getProjectId())
        .setRetrySettings(RETRY_SETTINGS).build().getService();
  }

  // Get and cache a listing of table names for this dataset.
  Set<String> tablesInDataset;
  if (tableListingCache == null) {
    // We need to be very careful about settings for the cache here. We have had significant
    // issues in the past due to exceeding limits on BigQuery API requests; see
    // https://bugzilla.mozilla.org/show_bug.cgi?id=1623000
    tableListingCache = CacheBuilder.newBuilder().expireAfterWrite(Duration.ofMinutes(10))
        .build();
  }
  try {
    tablesInDataset = tableListingCache.get(datasetRef, () -> {
      Set<String> tableSet = new HashSet<>();
      Dataset dataset = bqService.getDataset(ref.getDatasetId());
      if (dataset != null) {
        dataset.list().iterateAll().forEach(t -> {
          tableSet.add(t.getTableId().getTable());
        });
      }
      return tableSet;
    });
  } catch (ExecutionException e) {
    throw new UncheckedExecutionException(e.getCause());
  }

  // Send to error collection if dataset or table doesn't exist so BigQueryIO doesn't throw a
  // pipeline execution exception.
  if (tablesInDataset.isEmpty()) {
    throw new IllegalArgumentException("Resolved destination dataset does not exist or has no "
        + " tables for tableSpec " + tableSpec);
  } else if (!tablesInDataset.contains(ref.getTableId())) {
    throw new IllegalArgumentException("Resolved destination table does not exist: " + tableSpec);
  }

  return tableDestination;
}
 

private BigQueryRunner() throws IOException {
  this(
      MetricServiceClient.create(),
      BigQueryOptions.getDefaultInstance().getService(),
      System.out);
}
 

private static com.google.cloud.bigquery.BigQuery getBigQueryService(Env env) {
  return BigQueryOptions.getDefaultInstance().getService();
}
 

@Setup
public void setUp() {
  if (bigQueryClient == null) {
    bigQueryClient = BigQueryOptions.getDefaultInstance().getService();
  }
}
 

protected BigQuery authenticatedBQ() {
    return BigQueryOptions.newBuilder()
            .setCredentials(getGoogleCredentials())
            .build().getService();
}
 

/**
 * Creates the publisher with a application default credentials from the environment and default
 * project name.
 */
public static BigQueryClient create(String dataset) {
  BigQueryOptions options = BigQueryOptions.newBuilder().build();

  return new BigQueryClient(options.getService(), options.getProjectId(), dataset);
}