Skip to content

Commit

Permalink
Add helper method to create context
Browse files Browse the repository at this point in the history
  • Loading branch information
@yihanzhen
yihanzhen committed Apr 12, 2018
1 parent f539045 commit a6e6866
Showing 1 changed file with 32 additions and 53 deletions.
85 changes: 32 additions & 53 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,8 +174,7 @@ public Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable Str
}
ListInstanceConfigsRequest request = requestBuilder.build();

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
GrpcCallContext context = newCallContext(null, projectName);
return get(instanceStub.getInstanceConfigCallable().futureCall(request, context));
}

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

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(parent, projectName));
GrpcCallContext context = newCallContext(null, parent);
return get(instanceStub.createInstanceCallable().futureCall(request, context));
}

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instance.getName(), projectName));
GrpcCallContext context = newCallContext(null, instance.getName());
return get(instanceStub.updateInstanceCallable().futureCall(request, context));
}

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
GrpcCallContext context = newCallContext(null, instanceName);
return get(instanceStub.getInstanceCallable().futureCall(request, context));
}

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
GrpcCallContext context = newCallContext(null, instanceName);
get(instanceStub.deleteInstanceCallable().futureCall(request, context));
}

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
GrpcCallContext context = newCallContext(null, instanceName);
ListDatabasesResponse response = get(databaseStub.listDatabasesCallable()
.futureCall(request, context));
return new Paginated<>(response.getDatabasesList(), response.getNextPageToken());
Expand All @@ -282,8 +274,7 @@ public Operation createDatabase(String instanceName, String createDatabaseStatem
.setCreateStatement(createDatabaseStatement)
.addAllExtraStatements(additionalStatements)
.build();
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
GrpcCallContext context = newCallContext(null, instanceName);
return get(databaseStub.createDatabaseCallable().futureCall(request, context));
}

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

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
GrpcCallContext context = newCallContext(null, databaseName);
get(databaseStub.dropDatabaseCallable().futureCall(request, context));
}

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
GrpcCallContext context = newCallContext(null, databaseName);
return get(databaseStub.getDatabaseCallable().futureCall(request, context));
}

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

GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
GrpcCallContext context = newCallContext(null, databaseName);
return get(databaseStub.getDatabaseDdlCallable().futureCall(request, context))
.getStatementsList();
}

@Override
public Operation getOperation(String name) throws SpannerException {
GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build();
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(name, projectName));
GrpcCallContext context = newCallContext(null, projectName);
return get(databaseStub.getOperationsStub().getOperationCallable()
.futureCall(request, context));
}
Expand All @@ -353,9 +339,7 @@ public Session createSession(String databaseName, @Nullable Map<String, String>
requestBuilder.setSession(session);
}
CreateSessionRequest request = requestBuilder.build();
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
GrpcCallContext context = newCallContext(options, databaseName);
return get(stub.createSessionCallable().futureCall(request, context));
}

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

@Override
public StreamingCall read(
ReadRequest request, ResultStreamConsumer consumer, @Nullable Map<Option, ?> options) {
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
GrpcCallContext context = newCallContext(options, request.getSession());
throw new UnsupportedOperationException("not implemented yet");
}

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

@Override
public CommitResponse commit(CommitRequest commitRequest, @Nullable Map<Option, ?> options)
throws SpannerException {
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(
metadataProvider.newExtraHeaders(commitRequest.getSession(), projectName));
GrpcCallContext context = newCallContext(options, commitRequest.getSession());
return get(stub.commitCallable().futureCall(commitRequest, context));
}

@Override
public void rollback(RollbackRequest request, @Nullable Map<Option, ?> options)
throws SpannerException {
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
GrpcCallContext context = newCallContext(options, request.getSession());
get(stub.rollbackCallable().futureCall(request, context));
}

@Override
public PartitionResponse partitionQuery(
PartitionQueryRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
GrpcCallContext context = newCallContext(options, request.getSession());
return get(stub.partitionQueryCallable().futureCall(request, context));
}

@Override
public PartitionResponse partitionRead(
PartitionReadRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
GrpcCallContext context = newCallContext(options, request.getSession());
return get(stub.partitionReadCallable().futureCall(request, context));
}

Expand All @@ -444,4 +413,14 @@ private static <T> T get(final Future<T> future) throws SpannerException {
throw newSpannerException(context, e);
}
}

private GrpcCallContext newCallContext(@Nullable Map<Option, ?> options, String resource) {
GrpcCallContext context = GrpcCallContext.createDefault();
if (options != null) {
context = context.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue());
}
context = context.withExtraHeaders(
metadataProvider.newExtraHeaders(resource, projectName));
return context;
}
}

0 comments on commit a6e6866

Please sign in to comment.