Examples with BigQueryStorageClient com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient used on opensource projects

Example 1 with BigQueryStorageClient

use of com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient in project presto by prestodb.

the class ReadSessionCreator method create.

public Storage.ReadSession create(TableId table, ImmutableList<String> selectedFields, Optional<String> filter, int parallelism) {
    TableInfo tableDetails = bigQueryClient.getTable(table);
    TableInfo actualTable = getActualTable(tableDetails, selectedFields, new String[] {});
    try (BigQueryStorageClient bigQueryStorageClient = bigQueryStorageClientFactory.createBigQueryStorageClient()) {
        ReadOptions.TableReadOptions.Builder readOptions = ReadOptions.TableReadOptions.newBuilder().addAllSelectedFields(selectedFields);
        filter.ifPresent(readOptions::setRowRestriction);
        TableReferenceProto.TableReference tableReference = toTableReference(actualTable.getTableId());
        Storage.ReadSession readSession = bigQueryStorageClient.createReadSession(Storage.CreateReadSessionRequest.newBuilder().setParent("projects/" + bigQueryClient.getProjectId()).setFormat(Storage.DataFormat.AVRO).setRequestedStreams(parallelism).setReadOptions(readOptions).setTableReference(tableReference).setShardingStrategy(Storage.ShardingStrategy.BALANCED).build());
        return readSession;
    }
}

Example 2 with BigQueryStorageClient

use of com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient in project presto by prestodb.

the class BigQueryStorageClientFactory method createBigQueryStorageClient.

BigQueryStorageClient createBigQueryStorageClient() {
    try {
        BigQueryStorageSettings.Builder clientSettings = BigQueryStorageSettings.newBuilder().setTransportChannelProvider(BigQueryStorageSettings.defaultGrpcTransportProviderBuilder().setHeaderProvider(headerProvider).build());
        credentials.ifPresent(credentials -> clientSettings.setCredentialsProvider(FixedCredentialsProvider.create(credentials)));
        return BigQueryStorageClient.create(clientSettings.build());
    } catch (IOException e) {
        throw new UncheckedIOException("Error creating BigQueryStorageClient", e);
    }
}

Example 3 with BigQueryStorageClient

use of com.google.cloud.bigquery.storage.v1beta1.BigQueryStorageClient in project presto by prestodb.

the class ReadRowsHelper method readRows.

// TODO: iterator based implementation, instead of fetching all result at once
public Iterator<ReadRowsResponse> readRows() {
    List<ReadRowsResponse> readRowResponses = new ArrayList<>();
    long readRowsCount = 0;
    int retries = 0;
    Iterator<ReadRowsResponse> serverResponses = fetchResponses(request);
    while (serverResponses.hasNext()) {
        try {
            ReadRowsResponse response = serverResponses.next();
            readRowsCount += response.getRowCount();
            readRowResponses.add(response);
        } catch (RuntimeException e) {
            // if relevant, retry the read, from the last read position
            if (BigQueryUtil.isRetryable(e) && retries < maxReadRowsRetries) {
                request.getReadPositionBuilder().setOffset(readRowsCount);
                serverResponses = fetchResponses(request);
                retries++;
            } else {
                // to safely close the client
                try (BigQueryStorageClient ignored = client) {
                    throw e;
                }
            }
        }
    }
    return readRowResponses.iterator();
}