Exemplo n.º 1
0
  @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();
  }
Exemplo n.º 2
0
  /** 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;
  }
Exemplo n.º 3
0
 @After
 public void tearDown() throws Exception {
   final String tableName = bigtableOptions.getInstanceName().toTableNameStr(tableId);
   deleteTable(tableName);
   session.close();
 }