@VisibleForTesting
BigtableDataGrpcClient(
ChannelPool channelPool,
ScheduledExecutorService retryExecutorService,
BigtableOptions bigtableOptions,
BigtableAsyncUtilities asyncUtilities) {
this.channelPool = channelPool;
this.retryExecutorService = retryExecutorService;
this.bigtableOptions = bigtableOptions;
this.retryOptions = bigtableOptions.getRetryOptions();
this.sampleRowKeysAsync =
asyncUtilities.createAsyncRpc(
BigtableGrpc.METHOD_SAMPLE_ROW_KEYS, Predicates.<SampleRowKeysRequest>alwaysTrue());
this.readRowsAsync =
asyncUtilities.createAsyncRpc(
BigtableGrpc.METHOD_READ_ROWS, Predicates.<ReadRowsRequest>alwaysTrue());
this.mutateRowRpc =
asyncUtilities.createAsyncRpc(
BigtableGrpc.METHOD_MUTATE_ROW, getMutationRetryableFunction(IS_RETRYABLE_MUTATION));
this.mutateRowsRpc =
asyncUtilities.createAsyncRpc(
BigtableGrpc.METHOD_MUTATE_ROWS, getMutationRetryableFunction(ARE_RETRYABLE_MUTATIONS));
this.checkAndMutateRpc =
asyncUtilities.createAsyncRpc(
BigtableGrpc.METHOD_CHECK_AND_MUTATE_ROW,
getMutationRetryableFunction(IS_RETRYABLE_CHECK_AND_MUTATE));
this.readWriteModifyRpc =
asyncUtilities.createAsyncRpc(
BigtableGrpc.METHOD_READ_MODIFY_WRITE_ROW,
Predicates.<ReadModifyWriteRowRequest>alwaysFalse());
}
@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();
}
@Test
public void testE2EBigtableWrite() throws Exception {
final String tableName = bigtableOptions.getInstanceName().toTableNameStr(tableId);
final String instanceName = bigtableOptions.getInstanceName().toString();
final int numRows = 1000;
final List<KV<ByteString, ByteString>> testData = generateTableData(numRows);
createEmptyTable(instanceName, tableId);
Pipeline p = Pipeline.create(options);
p.apply(CountingInput.upTo(numRows))
.apply(
ParDo.of(
new DoFn<Long, KV<ByteString, Iterable<Mutation>>>() {
@ProcessElement
public void processElement(ProcessContext c) {
int index = c.element().intValue();
Iterable<Mutation> mutations =
ImmutableList.of(
Mutation.newBuilder()
.setSetCell(
Mutation.SetCell.newBuilder()
.setValue(testData.get(index).getValue())
.setFamilyName(COLUMN_FAMILY_NAME))
.build());
c.output(KV.of(testData.get(index).getKey(), mutations));
}
}))
.apply(BigtableIO.write().withBigtableOptions(bigtableOptions).withTableId(tableId));
p.run();
// Test number of column families and column family name equality
Table table = getTable(tableName);
assertThat(table.getColumnFamilies().keySet(), Matchers.hasSize(1));
assertThat(table.getColumnFamilies(), Matchers.hasKey(COLUMN_FAMILY_NAME));
// Test table data equality
List<KV<ByteString, ByteString>> tableData = getTableData(tableName);
assertThat(tableData, Matchers.containsInAnyOrder(testData.toArray()));
}
public static void main(String[] args) throws IOException {
Logger logger = new Logger(CheckConfig.class);
GenericOptionsParser optionsParser =
new GenericOptionsParser(HBaseConfiguration.create(), args);
Configuration fullConfiguration = optionsParser.getConfiguration();
BigtableOptions options;
try {
options = BigtableOptionsFactory.fromConfiguration(fullConfiguration);
} catch (IOException | RuntimeException exc) {
logger.warn("Encountered errors attempting to parse configuration.", exc);
return;
}
System.out.println(String.format("User Agent: %s", options.getChannelOptions().getUserAgent()));
System.out.println(String.format("Project ID: %s", options.getProjectId()));
System.out.println(String.format("Cluster Name: %s", options.getCluster()));
System.out.println(String.format("Zone: %s", options.getZone()));
System.out.println(String.format("Cluster admin host: %s", options.getClusterAdminHost()));
System.out.println(String.format("Table admin host: %s", options.getTableAdminHost()));
System.out.println(String.format("Data host: %s", options.getDataHost()));
Credentials credentials = options.getChannelOptions().getCredential();
try {
System.out.println("Attempting credential refresh...");
credentials.refresh();
} catch (IOException ioe) {
logger.warn("Encountered errors attempting to refresh credentials.", ioe);
return;
}
String configuredConnectionClass =
fullConfiguration.get(HConnection.HBASE_CLIENT_CONNECTION_IMPL);
boolean isCorrectClassSpecified = false;
if (!Strings.isNullOrEmpty(configuredConnectionClass)) {
try {
Class<?> connectionClass = Class.forName(configuredConnectionClass);
isCorrectClassSpecified =
AbstractBigtableConnection.class.isAssignableFrom(connectionClass);
} catch (Exception e) {
// Ignore. Problems will be logged in the println below.
}
}
// We can actually determine if this value is correct (disregarding custom subclasses).
System.out.println(
String.format(
"HBase Connection Class = %s %s",
configuredConnectionClass, isCorrectClassSpecified ? "(OK)" : "(Configuration error)"));
System.out.println("Opening table admin connection...");
try (Connection conn = ConnectionFactory.createConnection(fullConfiguration)) {
try (Admin admin = conn.getAdmin()) {
System.out.println(String.format("Tables in cluster %s:", options.getCluster()));
TableName[] tableNames = admin.listTableNames();
if (tableNames.length == 0) {
System.out.println("No tables found.");
} else {
for (TableName table : tableNames) {
System.out.println(table.getNameAsString());
}
}
}
System.out.println("Closing connection...");
}
}
@After
public void tearDown() throws Exception {
final String tableName = bigtableOptions.getInstanceName().toTableNameStr(tableId);
deleteTable(tableName);
session.close();
}