2727import com .google .cloud .bigtable .data .v2 .BigtableDataClient ;
2828import com .google .cloud .bigtable .data .v2 .BigtableDataSettings ;
2929import com .google .cloud .bigtable .data .v2 .stub .EnhancedBigtableStubSettings ;
30- import com .google .common .base .Joiner ;
3130import com .google .common .base .MoreObjects ;
3231import com .google .common .base .Predicate ;
3332import com .google .common .base .Predicates ;
4342import io .grpc .ManagedChannelBuilder ;
4443import io .grpc .Metadata ;
4544import io .grpc .MethodDescriptor ;
45+ import io .grpc .Status ;
4646import java .io .IOException ;
4747import java .net .InetSocketAddress ;
4848import java .net .SocketAddress ;
49- import java .util .ArrayList ;
5049import java .util .HashMap ;
51- import java .util .List ;
5250import java .util .Map ;
5351import java .util .Optional ;
5452import javax .annotation .Nullable ;
6361 * <li>{@code bigtable.table}
6462 * </ul>
6563 */
66- class CloudEnv extends AbstractTestEnv {
64+ public class CloudEnv extends AbstractTestEnv {
6765 private static final Predicate <InetSocketAddress > DIRECT_PATH_IPV6_MATCHER =
6866 new Predicate <InetSocketAddress >() {
6967 @ Override
@@ -84,6 +82,7 @@ public boolean apply(InetSocketAddress input) {
8482
8583 private static final String PROJECT_PROPERTY_NAME = "bigtable.project" ;
8684 private static final String INSTANCE_PROPERTY_NAME = "bigtable.instance" ;
85+ private static final String APP_PROFILE_PROPERTY_NAME = "bigtable.app_profile" ;
8786 private static final String TABLE_PROPERTY_NAME = "bigtable.table" ;
8887 private static final String CMEK_KMS_KEY_PROPERTY_NAME = "bigtable.kms_key_name" ;
8988
@@ -92,12 +91,12 @@ public boolean apply(InetSocketAddress input) {
9291 private final String projectId ;
9392 private final String instanceId ;
9493 private final String tableId ;
95- private final String tracingCookie ;
9694 private final String kmsKeyName ;
9795
9896 private final BigtableDataSettings .Builder dataSettings ;
9997 private final BigtableTableAdminSettings .Builder tableAdminSettings ;
10098 private final BigtableInstanceAdminSettings .Builder instanceAdminSettings ;
99+ @ Nullable private final String appProfileId ;
101100
102101 private BigtableDataClient dataClient ;
103102 private BigtableTableAdminClient tableAdminClient ;
@@ -110,6 +109,7 @@ static CloudEnv fromSystemProperties() {
110109 getOptionalProperty (CMEK_KMS_KEY_PROPERTY_NAME , "" ),
111110 getRequiredProperty (PROJECT_PROPERTY_NAME ),
112111 getRequiredProperty (INSTANCE_PROPERTY_NAME ),
112+ getOptionalProperty (APP_PROFILE_PROPERTY_NAME ),
113113 getRequiredProperty (TABLE_PROPERTY_NAME ),
114114 getOptionalProperty (TRACING_COOKIE_PROPERTY_NAME ));
115115 }
@@ -120,19 +120,23 @@ private CloudEnv(
120120 @ Nullable String kmsKeyName ,
121121 String projectId ,
122122 String instanceId ,
123+ @ Nullable String appProfileId ,
123124 String tableId ,
124125 @ Nullable String tracingCookie ) {
125126 this .projectId = projectId ;
126127 this .instanceId = instanceId ;
128+ this .appProfileId = appProfileId ;
127129 this .tableId = tableId ;
128- this .tracingCookie = tracingCookie ;
129130 this .kmsKeyName = kmsKeyName ;
130131
131132 this .dataSettings =
132133 BigtableDataSettings .newBuilder ().setProjectId (projectId ).setInstanceId (instanceId );
133134 if (!Strings .isNullOrEmpty (dataEndpoint )) {
134135 dataSettings .stubSettings ().setEndpoint (dataEndpoint );
135136 }
137+ if (!Strings .isNullOrEmpty (appProfileId )) {
138+ dataSettings .setAppProfileId (appProfileId );
139+ }
136140
137141 configureConnection (dataSettings .stubSettings ());
138142 configureUserAgent (dataSettings .stubSettings ());
@@ -193,6 +197,9 @@ private void configureConnection(StubSettings.Builder stubSettings) {
193197 throw new IllegalStateException ("Unexpected ConnectionMode: " + getConnectionMode ());
194198 }
195199
200+ final ClientInterceptor appProfileInterceptor =
201+ appProfileId != null ? new AppProfileInterceptor () : null ;
202+
196203 // Inject the interceptor into the channel provider, taking care to preserve existing channel
197204 // configurator
198205 InstantiatingGrpcChannelProvider .Builder channelProvider =
@@ -211,7 +218,11 @@ public ManagedChannelBuilder apply(ManagedChannelBuilder builder) {
211218 if (oldConfigurator != null ) {
212219 builder = oldConfigurator .apply (builder );
213220 }
214- return builder .intercept (interceptor );
221+ builder = builder .intercept (interceptor );
222+ if (appProfileInterceptor != null ) {
223+ builder = builder .intercept (appProfileInterceptor );
224+ }
225+ return builder ;
215226 }
216227 };
217228 channelProvider .setChannelConfigurator (newConfigurator );
@@ -255,25 +266,35 @@ public void onHeaders(Metadata headers) {
255266 };
256267 }
257268
258- private void configureUserAgent (EnhancedBigtableStubSettings .Builder stubSettings ) {
259- List <String > parts = new ArrayList <>();
260- parts .add ("java-bigtable-int-test" );
261-
262- switch (getConnectionMode ()) {
263- case DEFAULT :
264- // nothing special
265- break ;
266- case REQUIRE_CFE :
267- parts .add ("bigtable-directpath-disable" );
268- break ;
269- case REQUIRE_DIRECT_PATH :
270- case REQUIRE_DIRECT_PATH_IPV4 :
271- parts .add ("bigtable-directpath-enable" );
272- break ;
273- default :
274- throw new IllegalStateException ("Unexpected connectionMode: " + getConnectionMode ());
269+ private class AppProfileInterceptor implements ClientInterceptor {
270+ @ Override
271+ public <ReqT , RespT > ClientCall <ReqT , RespT > interceptCall (
272+ MethodDescriptor <ReqT , RespT > methodDescriptor , CallOptions callOptions , Channel channel ) {
273+ return new SimpleForwardingClientCall <ReqT , RespT >(
274+ channel .newCall (methodDescriptor , callOptions )) {
275+ @ Override
276+ public void start (Listener <RespT > responseListener , Metadata headers ) {
277+ String reqParams =
278+ headers .get (
279+ Metadata .Key .of ("x-goog-request-params" , Metadata .ASCII_STRING_MARSHALLER ));
280+ if (!reqParams .contains ("app_profile_id=" + appProfileId )) {
281+ responseListener .onClose (
282+ Status .FAILED_PRECONDITION .withDescription (
283+ "Integration test was configured to run with app profile: "
284+ + appProfileId
285+ + ", but found a different app profile in the headers: "
286+ + reqParams ),
287+ new Metadata ());
288+ return ;
289+ }
290+ super .start (responseListener , headers );
291+ }
292+ };
275293 }
276- String newUserAgent = Joiner .on (" " ).join (parts );
294+ }
295+
296+ private void configureUserAgent (EnhancedBigtableStubSettings .Builder stubSettings ) {
297+ String newUserAgent = "java-bigtable-int-test" ;
277298
278299 // Use the existing user-agent to use as a prefix
279300 Map <String , String > existingHeaders =
@@ -309,19 +330,6 @@ public BigtableDataClient getDataClient() {
309330 return dataClient ;
310331 }
311332
312- @ Override
313- public BigtableDataClient getDataClientForInstance (String instanceId ) throws IOException {
314- BigtableDataSettings .Builder settings =
315- BigtableDataSettings .newBuilder ()
316- .setProjectId (dataSettings .getProjectId ())
317- .setInstanceId (instanceId );
318- settings
319- .stubSettings ()
320- .setEndpoint (dataSettings .stubSettings ().getEndpoint ())
321- .setTransportChannelProvider (dataSettings .stubSettings ().getTransportChannelProvider ());
322- return BigtableDataClient .create (settings .build ());
323- }
324-
325333 @ Override
326334 public BigtableTableAdminClient getTableAdminClient () {
327335 return tableAdminClient ;
0 commit comments