diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java index 5f621654a..4b5ede9da 100644 --- a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java @@ -238,7 +238,9 @@ private ServerStreamingCallable createReadRowsBaseCallable( new RequestParamsExtractor() { @Override public Map extract(ReadRowsRequest readRowsRequest) { - return ImmutableMap.of("table_name", readRowsRequest.getTableName()); + return ImmutableMap.of( + "table_name", readRowsRequest.getTableName(), + "app_profile_id", readRowsRequest.getAppProfileId()); } }) .build(), @@ -296,7 +298,9 @@ private UnaryCallable> createSampleRowKeysCallable() { @Override public Map extract( SampleRowKeysRequest sampleRowKeysRequest) { - return ImmutableMap.of("table_name", sampleRowKeysRequest.getTableName()); + return ImmutableMap.of( + "table_name", sampleRowKeysRequest.getTableName(), + "app_profile_id", sampleRowKeysRequest.getAppProfileId()); } }) .build(), @@ -328,7 +332,9 @@ private UnaryCallable createMutateRowCallable() { new RequestParamsExtractor() { @Override public Map extract(MutateRowRequest mutateRowRequest) { - return ImmutableMap.of("table_name", mutateRowRequest.getTableName()); + return ImmutableMap.of( + "table_name", mutateRowRequest.getTableName(), + "app_profile_id", mutateRowRequest.getAppProfileId()); } }) .build(), @@ -449,7 +455,9 @@ private UnaryCallable createMutateRowsBaseCallable() { new RequestParamsExtractor() { @Override public Map extract(MutateRowsRequest mutateRowsRequest) { - return ImmutableMap.of("table_name", mutateRowsRequest.getTableName()); + return ImmutableMap.of( + "table_name", mutateRowsRequest.getTableName(), + "app_profile_id", mutateRowsRequest.getAppProfileId()); } }) .build(), @@ -490,7 +498,8 @@ private UnaryCallable createCheckAndMutateRowCa public Map extract( CheckAndMutateRowRequest checkAndMutateRowRequest) { return ImmutableMap.of( - "table_name", checkAndMutateRowRequest.getTableName()); + "table_name", checkAndMutateRowRequest.getTableName(), + "app_profile_id", checkAndMutateRowRequest.getAppProfileId()); } }) .build(), @@ -522,7 +531,9 @@ private UnaryCallable createReadModifyWriteRowCallable( new RequestParamsExtractor() { @Override public Map extract(ReadModifyWriteRowRequest request) { - return ImmutableMap.of("table_name", request.getTableName()); + return ImmutableMap.of( + "table_name", request.getTableName(), + "app_profile_id", request.getAppProfileId()); } }) .build(), diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ResourceHeaderTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ResourceHeaderTest.java index b65c92ed8..85c8254ce 100644 --- a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ResourceHeaderTest.java +++ b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/ResourceHeaderTest.java @@ -44,13 +44,14 @@ public class ResourceHeaderTest { private static final String PROJECT_ID = "fake-project"; private static final String INSTANCE_ID = "fake-instance"; private static final String TABLE_ID = "fake-table"; - private static final String NAME = "resource-header-test:123"; - private static final Pattern EXPECTED_HEADER_PATTERN = - Pattern.compile(".*" + NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID) + ".*"); - private static final String HEADER_NAME = "x-goog-request-params"; - private static final String TEST_HEADER_NAME = "simple-header-name"; + private static final String TABLE_NAME = + NameUtil.formatTableName(PROJECT_ID, INSTANCE_ID, TABLE_ID); + private static final String APP_PROFILE_ID = "fake-profile"; + + private static final String CHANNEL_NAME = "resource-header-test:123"; + private static final String X_GOOG_REQUEST_PARAMS_KEY = "x-goog-request-params"; + private static final String TEST_HEADER_KEY = "simple-header-name"; private static final String TEST_HEADER_VALUE = "simple-header-value"; - private static final Pattern TEST_PATTERN = Pattern.compile(".*" + TEST_HEADER_VALUE + ".*"); private InProcessServer server; private LocalChannelProvider channelProvider; @@ -58,20 +59,21 @@ public class ResourceHeaderTest { @Before public void setUp() throws Exception { - server = new InProcessServer<>(new BigtableGrpc.BigtableImplBase() {}, NAME); + server = new InProcessServer<>(new BigtableGrpc.BigtableImplBase() {}, CHANNEL_NAME); server.start(); - channelProvider = LocalChannelProvider.create(NAME); + channelProvider = LocalChannelProvider.create(CHANNEL_NAME); BigtableDataSettings.Builder settings = BigtableDataSettings.newBuilder() .setProjectId(PROJECT_ID) .setInstanceId(INSTANCE_ID) + .setAppProfileId(APP_PROFILE_ID) .setCredentialsProvider(NoCredentialsProvider.create()); settings .stubSettings() .setTransportChannelProvider(channelProvider) - .setHeaderProvider(FixedHeaderProvider.create(TEST_HEADER_NAME, TEST_HEADER_VALUE)); + .setHeaderProvider(FixedHeaderProvider.create(TEST_HEADER_KEY, TEST_HEADER_VALUE)); // Force immediate flush settings @@ -139,9 +141,21 @@ public void readModifyWriteTest() { } private void verifyHeaderSent() { - boolean headerSent = channelProvider.isHeaderSent(HEADER_NAME, EXPECTED_HEADER_PATTERN); - assertWithMessage("Header was sent").that(headerSent).isTrue(); - boolean testHeader = channelProvider.isHeaderSent(TEST_HEADER_NAME, TEST_PATTERN); + boolean tableHeaderSent = + channelProvider.isHeaderSent( + X_GOOG_REQUEST_PARAMS_KEY, + Pattern.compile("(^|.*&)table_name=" + TABLE_NAME + "($|&.*)")); + assertWithMessage("Tablename header was sent").that(tableHeaderSent).isTrue(); + + boolean appProfileHeaderSent = + channelProvider.isHeaderSent( + X_GOOG_REQUEST_PARAMS_KEY, + Pattern.compile("(^|.*&)app_profile_id=" + APP_PROFILE_ID + "($|&.*)")); + assertWithMessage("App profile header was sent").that(appProfileHeaderSent).isTrue(); + + boolean testHeader = + channelProvider.isHeaderSent( + TEST_HEADER_KEY, Pattern.compile("^" + TEST_HEADER_VALUE + "$")); assertWithMessage("HeaderProvider's header received in Channel").that(testHeader).isTrue(); } }