diff --git a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java index 5ac1a3c775..3a696dcfe7 100644 --- a/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java +++ b/test-proxy/src/main/java/com/google/cloud/bigtable/testproxy/CbtTestProxy.java @@ -202,7 +202,7 @@ public synchronized void createClient( BigtableDataSettings.Builder settingsBuilder = BigtableDataSettings.newBuilder() - // disable channel refreshing when creating an emulator + // Disable channel refreshing when not using the real server .setRefreshingChannel(false) .setProjectId(request.getProjectId()) .setInstanceId(request.getInstanceId()) @@ -217,6 +217,19 @@ public synchronized void createClient( Durations.toString(request.getPerOperationTimeout()))); } + if (request.getOptionalFeatureConfig() + == OptionalFeatureConfig.OPTIONAL_FEATURE_CONFIG_ENABLE_ALL) { + logger.info("Enabling all the optional features"); + try { + // Exception will be raised if Application Default Credentials is not found. + BigtableDataSettings.enableBuiltinMetrics(); + } catch (IOException e) { + responseObserver.onError(Status.INTERNAL.withDescription(e.getMessage()).asException()); + return; + } + settingsBuilder.stubSettings().bulkMutateRowsSettings().setServerInitiatedFlowControl(true); + } + // Create and store CbtClient for later use try { if (!request.getDataTarget().equals("emulator")) { diff --git a/test-proxy/src/main/proto/test_proxy.proto b/test-proxy/src/main/proto/test_proxy.proto index 551dd4d8b2..e7caef0e7b 100644 --- a/test-proxy/src/main/proto/test_proxy.proto +++ b/test-proxy/src/main/proto/test_proxy.proto @@ -26,6 +26,16 @@ option go_package = "./testproxypb"; option java_multiple_files = true; option java_package = "com.google.cloud.bigtable.testproxy"; +// A config flag that dictates how the optional features should be enabled +// during the client creation. The optional features customize how the client +// interacts with the server, and are defined in +// https://github.com/googleapis/googleapis/blob/master/google/bigtable/v2/feature_flags.proto +enum OptionalFeatureConfig { + OPTIONAL_FEATURE_CONFIG_DEFAULT = 0; + + OPTIONAL_FEATURE_CONFIG_ENABLE_ALL = 1; +} + // Request to test proxy service to create a client object. message CreateClientRequest { // A unique ID associated with the client object to be created. @@ -52,6 +62,10 @@ message CreateClientRequest { // the created client. Otherwise, the default timeout from the client library // will be used. Note that the override applies to all the methods. google.protobuf.Duration per_operation_timeout = 6; + + // Optional config that dictates how the optional features should be enabled + // during the client creation. Please check the enum type's docstring above. + OptionalFeatureConfig optional_feature_config = 7; } // Response from test proxy service for CreateClientRequest.