Skip to content

Commit aecf956

Browse files
yihanzhenyihanzhen
committed
Add helper method to create context
1 parent f539045 commit aecf956

File tree

1 file changed

+32
-53
lines changed

1 file changed

+32
-53
lines changed

‎google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java‎

Lines changed: 32 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -174,8 +174,7 @@ public Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable Str
174174
}
175175
ListInstanceConfigsRequest request = requestBuilder.build();
176176

177-
GrpcCallContext context = GrpcCallContext.createDefault()
178-
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
177+
GrpcCallContext context = newCallContext(null, projectName);
179178
ListInstanceConfigsResponse response =
180179
get(instanceStub.listInstanceConfigsCallable().futureCall(request, context));
181180
return new Paginated<>(response.getInstanceConfigsList(), response.getNextPageToken());
@@ -186,8 +185,7 @@ public InstanceConfig getInstanceConfig(String instanceConfigName) throws Spanne
186185
GetInstanceConfigRequest request =
187186
GetInstanceConfigRequest.newBuilder().setName(instanceConfigName).build();
188187

189-
GrpcCallContext context = GrpcCallContext.createDefault()
190-
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
188+
GrpcCallContext context = newCallContext(null, projectName);
191189
return get(instanceStub.getInstanceConfigCallable().futureCall(request, context));
192190
}
193191

@@ -204,8 +202,7 @@ public Paginated<Instance> listInstances(
204202
}
205203
ListInstancesRequest request = requestBuilder.build();
206204

207-
GrpcCallContext context = GrpcCallContext.createDefault()
208-
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
205+
GrpcCallContext context = newCallContext(null, projectName);
209206
ListInstancesResponse response =
210207
get(instanceStub.listInstancesCallable().futureCall(request, context));
211208
return new Paginated<>(response.getInstancesList(), response.getNextPageToken());
@@ -221,8 +218,7 @@ public Operation createInstance(String parent, String instanceId, Instance insta
221218
.setInstance(instance)
222219
.build();
223220

224-
GrpcCallContext context = GrpcCallContext.createDefault()
225-
.withExtraHeaders(metadataProvider.newExtraHeaders(parent, projectName));
221+
GrpcCallContext context = newCallContext(null, parent);
226222
return get(instanceStub.createInstanceCallable().futureCall(request, context));
227223
}
228224

@@ -231,8 +227,7 @@ public Operation updateInstance(Instance instance, FieldMask fieldMask) throws S
231227
UpdateInstanceRequest request =
232228
UpdateInstanceRequest.newBuilder().setInstance(instance).setFieldMask(fieldMask).build();
233229

234-
GrpcCallContext context = GrpcCallContext.createDefault()
235-
.withExtraHeaders(metadataProvider.newExtraHeaders(instance.getName(), projectName));
230+
GrpcCallContext context = newCallContext(null, instance.getName());
236231
return get(instanceStub.updateInstanceCallable().futureCall(request, context));
237232
}
238233

@@ -241,8 +236,7 @@ public Instance getInstance(String instanceName) throws SpannerException {
241236
GetInstanceRequest request =
242237
GetInstanceRequest.newBuilder().setName(instanceName).build();
243238

244-
GrpcCallContext context = GrpcCallContext.createDefault()
245-
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
239+
GrpcCallContext context = newCallContext(null, instanceName);
246240
return get(instanceStub.getInstanceCallable().futureCall(request, context));
247241
}
248242

@@ -251,8 +245,7 @@ public void deleteInstance(String instanceName) throws SpannerException {
251245
DeleteInstanceRequest request =
252246
DeleteInstanceRequest.newBuilder().setName(instanceName).build();
253247

254-
GrpcCallContext context = GrpcCallContext.createDefault()
255-
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
248+
GrpcCallContext context = newCallContext(null, instanceName);
256249
get(instanceStub.deleteInstanceCallable().futureCall(request, context));
257250
}
258251

@@ -266,8 +259,7 @@ public Paginated<Database> listDatabases(
266259
}
267260
ListDatabasesRequest request = requestBuilder.build();
268261

269-
GrpcCallContext context = GrpcCallContext.createDefault()
270-
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
262+
GrpcCallContext context = newCallContext(null, instanceName);
271263
ListDatabasesResponse response = get(databaseStub.listDatabasesCallable()
272264
.futureCall(request, context));
273265
return new Paginated<>(response.getDatabasesList(), response.getNextPageToken());
@@ -282,8 +274,7 @@ public Operation createDatabase(String instanceName, String createDatabaseStatem
282274
.setCreateStatement(createDatabaseStatement)
283275
.addAllExtraStatements(additionalStatements)
284276
.build();
285-
GrpcCallContext context = GrpcCallContext.createDefault()
286-
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
277+
GrpcCallContext context = newCallContext(null, instanceName);
287278
return get(databaseStub.createDatabaseCallable().futureCall(request, context));
288279
}
289280

@@ -296,8 +287,7 @@ public Operation updateDatabaseDdl(String databaseName, Iterable<String> updateD
296287
.addAllStatements(updateDatabaseStatements)
297288
.setOperationId(MoreObjects.firstNonNull(updateId, ""))
298289
.build();
299-
GrpcCallContext context = GrpcCallContext.createDefault()
300-
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
290+
GrpcCallContext context = newCallContext(null, databaseName);
301291
return get(databaseStub.updateDatabaseDdlCallable().futureCall(request, context));
302292
}
303293

@@ -306,8 +296,7 @@ public void dropDatabase(String databaseName) throws SpannerException {
306296
DropDatabaseRequest request =
307297
DropDatabaseRequest.newBuilder().setDatabase(databaseName).build();
308298

309-
GrpcCallContext context = GrpcCallContext.createDefault()
310-
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
299+
GrpcCallContext context = newCallContext(null, databaseName);
311300
get(databaseStub.dropDatabaseCallable().futureCall(request, context));
312301
}
313302

@@ -318,8 +307,7 @@ public Database getDatabase(String databaseName) throws SpannerException {
318307
.setName(databaseName)
319308
.build();
320309

321-
GrpcCallContext context = GrpcCallContext.createDefault()
322-
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
310+
GrpcCallContext context = newCallContext(null, databaseName);
323311
return get(databaseStub.getDatabaseCallable().futureCall(request, context));
324312
}
325313

@@ -328,17 +316,15 @@ public List<String> getDatabaseDdl(String databaseName) throws SpannerException
328316
GetDatabaseDdlRequest request =
329317
GetDatabaseDdlRequest.newBuilder().setDatabase(databaseName).build();
330318

331-
GrpcCallContext context = GrpcCallContext.createDefault()
332-
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
319+
GrpcCallContext context = newCallContext(null, databaseName);
333320
return get(databaseStub.getDatabaseDdlCallable().futureCall(request, context))
334321
.getStatementsList();
335322
}
336323

337324
@Override
338325
public Operation getOperation(String name) throws SpannerException {
339326
GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build();
340-
GrpcCallContext context = GrpcCallContext.createDefault()
341-
.withExtraHeaders(metadataProvider.newExtraHeaders(name, projectName));
327+
GrpcCallContext context = newCallContext(null, name);
342328
return get(databaseStub.getOperationsStub().getOperationCallable()
343329
.futureCall(request, context));
344330
}
@@ -353,9 +339,7 @@ public Session createSession(String databaseName, @Nullable Map<String, String>
353339
requestBuilder.setSession(session);
354340
}
355341
CreateSessionRequest request = requestBuilder.build();
356-
GrpcCallContext context = GrpcCallContext.createDefault()
357-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
358-
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
342+
GrpcCallContext context = newCallContext(options, databaseName);
359343
return get(stub.createSessionCallable().futureCall(request, context));
360344
}
361345

@@ -364,18 +348,14 @@ public void deleteSession(String sessionName, @Nullable Map<Option, ?> options)
364348
throws SpannerException {
365349
DeleteSessionRequest request =
366350
DeleteSessionRequest.newBuilder().setName(sessionName).build();
367-
GrpcCallContext context = GrpcCallContext.createDefault()
368-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
369-
.withExtraHeaders(metadataProvider.newExtraHeaders(sessionName, projectName));
351+
GrpcCallContext context = newCallContext(options, sessionName);
370352
get(stub.deleteSessionCallable().futureCall(request, context));
371353
}
372354

373355
@Override
374356
public StreamingCall read(
375357
ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
376-
GrpcCallContext context = GrpcCallContext.createDefault()
377-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
378-
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
358+
GrpcCallContext context = newCallContext(options, request.getSession());
379359
throw new UnsupportedOperationException("not implemented yet");
380360
}
381361

@@ -388,46 +368,35 @@ public StreamingCall executeQuery(
388368
@Override
389369
public Transaction beginTransaction(
390370
BeginTransactionRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
391-
GrpcCallContext context = GrpcCallContext.createDefault()
392-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
393-
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
371+
GrpcCallContext context = newCallContext(options, request.getSession());
394372
return get(stub.beginTransactionCallable().futureCall(request, context));
395373
}
396374

397375
@Override
398376
public CommitResponse commit(CommitRequest commitRequest, @Nullable Map<Option, ?> options)
399377
throws SpannerException {
400-
GrpcCallContext context = GrpcCallContext.createDefault()
401-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
402-
.withExtraHeaders(
403-
metadataProvider.newExtraHeaders(commitRequest.getSession(), projectName));
378+
GrpcCallContext context = newCallContext(options, commitRequest.getSession());
404379
return get(stub.commitCallable().futureCall(commitRequest, context));
405380
}
406381

407382
@Override
408383
public void rollback(RollbackRequest request, @Nullable Map<Option, ?> options)
409384
throws SpannerException {
410-
GrpcCallContext context = GrpcCallContext.createDefault()
411-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
412-
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
385+
GrpcCallContext context = newCallContext(options, request.getSession());
413386
get(stub.rollbackCallable().futureCall(request, context));
414387
}
415388

416389
@Override
417390
public PartitionResponse partitionQuery(
418391
PartitionQueryRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
419-
GrpcCallContext context = GrpcCallContext.createDefault()
420-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
421-
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
392+
GrpcCallContext context = newCallContext(options, request.getSession());
422393
return get(stub.partitionQueryCallable().futureCall(request, context));
423394
}
424395

425396
@Override
426397
public PartitionResponse partitionRead(
427398
PartitionReadRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
428-
GrpcCallContext context = GrpcCallContext.createDefault()
429-
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
430-
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
399+
GrpcCallContext context = newCallContext(options, request.getSession());
431400
return get(stub.partitionReadCallable().futureCall(request, context));
432401
}
433402

@@ -444,4 +413,14 @@ private static <T> T get(final Future<T> future) throws SpannerException {
444413
throw newSpannerException(context, e);
445414
}
446415
}
416+
417+
private GrpcCallContext newCallContext(@Nullable Map<Option, ?> options, String resource) {
418+
GrpcCallContext context = GrpcCallContext.createDefault();
419+
if (options != null) {
420+
context = context.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
421+
}
422+
context = context.withExtraHeaders(
423+
metadataProvider.newExtraHeaders(resource, projectName));
424+
return context;
425+
}
447426
}

0 commit comments

Comments
 (0)