Examples with BigQueryWriteClient com.google.cloud.bigquery.storage.v1.BigQueryWriteClient used on opensource projects

Example 1 with BigQueryWriteClient

use of com.google.cloud.bigquery.storage.v1.BigQueryWriteClient in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryClientFactoryTest method testGetWriteClientWithUserAgent.

@Test
public void testGetWriteClientWithUserAgent() {
    BigQueryClientFactory clientFactory = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, bigQueryConfig);
    when(bigQueryConfig.getBigQueryProxyConfig()).thenReturn(bigQueryProxyConfig);
    BigQueryWriteClient writeClient = clientFactory.getBigQueryWriteClient();
    assertNotNull(writeClient);
    BigQueryClientFactory clientFactory2 = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, bigQueryConfig);
    when(bigQueryConfig.getBigQueryProxyConfig()).thenReturn(bigQueryProxyConfig);
    BigQueryWriteClient writeClient2 = clientFactory2.getBigQueryWriteClient();
    assertNotNull(writeClient2);
    assertSame(writeClient, writeClient2);
    BigQueryClientFactory clientFactory3 = new BigQueryClientFactory(bigQueryCredentialsSupplier, HttpUtil.createHeaderProvider(bigQueryConfig, "test-agent-2"), bigQueryConfig);
    when(bigQueryConfig.getBigQueryProxyConfig()).thenReturn(bigQueryProxyConfig);
    BigQueryWriteClient writeClient3 = clientFactory3.getBigQueryWriteClient();
    assertNotNull(writeClient3);
    assertNotSame(writeClient, writeClient3);
    assertNotSame(writeClient2, writeClient3);
}

Example 2 with BigQueryWriteClient

use of com.google.cloud.bigquery.storage.v1.BigQueryWriteClient in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryClientFactoryTest method testGetWriteClientWithBigQueryConfig.

@Test
public void testGetWriteClientWithBigQueryConfig() {
    BigQueryClientFactory clientFactory = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, new TestBigQueryConfig(Optional.of("US:8080")));
    BigQueryWriteClient writeClient = clientFactory.getBigQueryWriteClient();
    assertNotNull(writeClient);
    BigQueryClientFactory clientFactory2 = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, new TestBigQueryConfig(Optional.of("US:8080")));
    BigQueryWriteClient writeClient2 = clientFactory2.getBigQueryWriteClient();
    assertNotNull(writeClient2);
    assertSame(writeClient, writeClient2);
    BigQueryClientFactory clientFactory3 = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, new TestBigQueryConfig(Optional.of("EU:8080")));
    BigQueryWriteClient writeClient3 = clientFactory3.getBigQueryWriteClient();
    assertNotNull(writeClient3);
    assertNotSame(writeClient, writeClient3);
    assertNotSame(writeClient2, writeClient3);
}

Example 3 with BigQueryWriteClient

use of com.google.cloud.bigquery.storage.v1.BigQueryWriteClient in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryClientFactoryTest method testGetWriteClientWithServiceAccountCredentials.

@Test
public void testGetWriteClientWithServiceAccountCredentials() throws Exception {
    when(bigQueryCredentialsSupplier.getCredentials()).thenReturn(createServiceAccountCredentials("test-client-id"));
    BigQueryClientFactory clientFactory = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, bigQueryConfig);
    when(bigQueryConfig.getBigQueryProxyConfig()).thenReturn(bigQueryProxyConfig);
    BigQueryWriteClient writeClient = clientFactory.getBigQueryWriteClient();
    assertNotNull(writeClient);
    when(bigQueryCredentialsSupplier.getCredentials()).thenReturn(createServiceAccountCredentials("test-client-id"));
    BigQueryClientFactory clientFactory2 = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, bigQueryConfig);
    when(bigQueryConfig.getBigQueryProxyConfig()).thenReturn(bigQueryProxyConfig);
    BigQueryWriteClient writeClient2 = clientFactory2.getBigQueryWriteClient();
    assertNotNull(writeClient2);
    assertSame(writeClient, writeClient2);
    when(bigQueryCredentialsSupplier.getCredentials()).thenReturn(createServiceAccountCredentials("test-client-id-2"));
    BigQueryClientFactory clientFactory3 = new BigQueryClientFactory(bigQueryCredentialsSupplier, headerProvider, bigQueryConfig);
    when(bigQueryConfig.getBigQueryProxyConfig()).thenReturn(bigQueryProxyConfig);
    BigQueryWriteClient writeClient3 = clientFactory3.getBigQueryWriteClient();
    assertNotNull(writeClient3);
    assertNotSame(writeClient, writeClient3);
    assertNotSame(writeClient2, writeClient3);
}

Example 4 with BigQueryWriteClient

use of com.google.cloud.bigquery.storage.v1.BigQueryWriteClient in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryClientFactory method getBigQueryWriteClient.

public BigQueryWriteClient getBigQueryWriteClient() {
    synchronized (writeClientMap) {
        if (!writeClientMap.containsKey(this)) {
            BigQueryWriteClient bigQueryWriteClient = createBigQueryWriteClient();
            Runtime.getRuntime().addShutdownHook(new Thread(() -> shutdownBigQueryWriteClient(bigQueryWriteClient)));
            writeClientMap.put(this, bigQueryWriteClient);
        }
    }
    return writeClientMap.get(this);
}

Example 5 with BigQueryWriteClient

use of com.google.cloud.bigquery.storage.v1.BigQueryWriteClient in project java-bigquerystorage by googleapis.

the class ParallelWriteCommittedStream method createStream.

private WriteStream createStream(String projectId, String datasetName, String tableName, BigQueryWriteClient client) {
    LOG.info("Creating a new stream");
    // Initialize a write stream for the specified table.
    // For more information on WriteStream.Type, see:
    // https://googleapis.dev/java/google-cloud-bigquerystorage/latest/com/google/cloud/bigquery/storage/v1/WriteStream.Type.html
    WriteStream stream = WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build();
    TableName parentTable = TableName.of(projectId, datasetName, tableName);
    CreateWriteStreamRequest createWriteStreamRequest = CreateWriteStreamRequest.newBuilder().setParent(parentTable.toString()).setWriteStream(stream).build();
    return client.createWriteStream(createWriteStreamRequest);
}