@Before
public void setup() throws Exception {
PipelineOptionsFactory.register(BigtableTestOptions.class);
options = TestPipeline.testingPipelineOptions().as(BigtableTestOptions.class);
// RetryOptions streamingBatchSize must be explicitly set for getTableData()
RetryOptions.Builder retryOptionsBuilder = new RetryOptions.Builder();
retryOptionsBuilder.setStreamingBatchSize(
retryOptionsBuilder.build().getStreamingBufferSize() / 2);
bigtableOptions =
new Builder()
.setProjectId(options.getProjectId())
.setInstanceId(options.getInstanceId())
.setUserAgent("apache-beam-test")
.setRetryOptions(retryOptionsBuilder.build())
.build();
session =
new BigtableSession(
bigtableOptions
.toBuilder()
.setCredentialOptions(
CredentialOptions.credential(options.as(GcpOptions.class).getGcpCredential()))
.build());
tableAdminClient = session.getTableAdminClient();
}
/** Helper function to get a table's data. */
private List<KV<ByteString, ByteString>> getTableData(String tableName) throws IOException {
// Add empty range to avoid TARGET_NOT_SET error
RowRange range =
RowRange.newBuilder()
.setStartKeyClosed(ByteString.EMPTY)
.setEndKeyOpen(ByteString.EMPTY)
.build();
RowSet rowSet = RowSet.newBuilder().addRowRanges(range).build();
ReadRowsRequest.Builder readRowsRequestBuilder =
ReadRowsRequest.newBuilder().setTableName(tableName).setRows(rowSet);
ResultScanner<Row> scanner = session.getDataClient().readRows(readRowsRequestBuilder.build());
Row currentRow;
List<KV<ByteString, ByteString>> tableData = new ArrayList<>();
while ((currentRow = scanner.next()) != null) {
ByteString key = currentRow.getKey();
ByteString value = currentRow.getFamilies(0).getColumns(0).getCells(0).getValue();
tableData.add(KV.of(key, value));
}
scanner.close();
return tableData;
}
@After
public void tearDown() throws Exception {
final String tableName = bigtableOptions.getInstanceName().toTableNameStr(tableId);
deleteTable(tableName);
session.close();
}
