Skip to content

Commit

Permalink
WIP: inject headers to grpcCallContext before making calls
Browse files Browse the repository at this point in the history
  • Loading branch information
@yihanzhen
yihanzhen committed Mar 30, 2018
1 parent ba0cf43 commit 9df854c
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 34 deletions.
82 changes: 51 additions & 31 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 @@ -21,6 +21,7 @@
import com.google.api.gax.core.CredentialsProvider;
import com.google.api.gax.core.GaxProperties;
import com.google.api.gax.grpc.GaxGrpcProperties;
import com.google.api.gax.grpc.GrpcCallContext;
import com.google.api.gax.grpc.GrpcTransportChannel;
import com.google.api.gax.rpc.ApiClientHeaderProvider;
import com.google.api.gax.rpc.FixedTransportChannelProvider;
Expand Down Expand Up @@ -163,9 +164,10 @@ public Paginated<InstanceConfig> listInstanceConfigs(int pageSize, @Nullable Str
}
ListInstanceConfigsRequest request = requestBuilder.build();

// TODO: put projectName in metadata
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
ListInstanceConfigsResponse response =
get(instanceStub.listInstanceConfigsCallable().futureCall(request));
get(instanceStub.listInstanceConfigsCallable().futureCall(request, context));
return new Paginated<>(response.getInstanceConfigsList(), response.getNextPageToken());
}

Expand All @@ -174,8 +176,9 @@ public InstanceConfig getInstanceConfig(String instanceConfigName) throws Spanne
GetInstanceConfigRequest request =
GetInstanceConfigRequest.newBuilder().setName(instanceConfigName).build();

// TODO: put projectName in metadata
return get(instanceStub.getInstanceConfigCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
return get(instanceStub.getInstanceConfigCallable().futureCall(request, context));
}

@Override
Expand All @@ -191,9 +194,10 @@ public Paginated<Instance> listInstances(
}
ListInstancesRequest request = requestBuilder.build();

// TODO: put projectName in metadata
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(projectName, projectName));
ListInstancesResponse response =
get(instanceStub.listInstancesCallable().futureCall(request));
get(instanceStub.listInstancesCallable().futureCall(request, context));
return new Paginated<>(response.getInstancesList(), response.getNextPageToken());
}

Expand All @@ -206,34 +210,40 @@ public Operation createInstance(String parent, String instanceId, Instance insta
.setInstanceId(instanceId)
.setInstance(instance)
.build();
// TODO: put parent in metadata
return get(instanceStub.createInstanceCallable().futureCall(request));

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

@Override
public Operation updateInstance(Instance instance, FieldMask fieldMask) throws SpannerException {
UpdateInstanceRequest request =
UpdateInstanceRequest.newBuilder().setInstance(instance).setFieldMask(fieldMask).build();
// TODO: put instance.getName() in metadata
return get(instanceStub.updateInstanceCallable().futureCall(request));

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

@Override
public Instance getInstance(String instanceName) throws SpannerException {
GetInstanceRequest request =
GetInstanceRequest.newBuilder().setName(instanceName).build();

// TODO: put instanceName in metadata
return get(instanceStub.getInstanceCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
return get(instanceStub.getInstanceCallable().futureCall(request, context));
}

@Override
public void deleteInstance(String instanceName) throws SpannerException {
DeleteInstanceRequest request =
DeleteInstanceRequest.newBuilder().setName(instanceName).build();

// TODO: put instanceName in metadata
get(instanceStub.deleteInstanceCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
get(instanceStub.deleteInstanceCallable().futureCall(request, context));
}

@Override
Expand All @@ -246,8 +256,10 @@ public Paginated<Database> listDatabases(
}
ListDatabasesRequest request = requestBuilder.build();

// TODO: put instanceName in metadata
ListDatabasesResponse response = get(databaseStub.listDatabasesCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
ListDatabasesResponse response = get(databaseStub.listDatabasesCallable()
.futureCall(request, context));
return new Paginated<>(response.getDatabasesList(), response.getNextPageToken());
}

Expand All @@ -260,8 +272,9 @@ public Operation createDatabase(String instanceName, String createDatabaseStatem
.setCreateStatement(createDatabaseStatement)
.addAllExtraStatements(additionalStatements)
.build();
// TODO: put instanceName in metadata
return get(databaseStub.createDatabaseCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(instanceName, projectName));
return get(databaseStub.createDatabaseCallable().futureCall(request, context));
}

@Override
Expand All @@ -273,17 +286,19 @@ public Operation updateDatabaseDdl(String databaseName, Iterable<String> updateD
.addAllStatements(updateDatabaseStatements)
.setOperationId(MoreObjects.firstNonNull(updateId, ""))
.build();
// TODO: put databaseName in metadata
return get(databaseStub.updateDatabaseDdlCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
return get(databaseStub.updateDatabaseDdlCallable().futureCall(request, context));
}

@Override
public void dropDatabase(String databaseName) throws SpannerException {
DropDatabaseRequest request =
DropDatabaseRequest.newBuilder().setDatabase(databaseName).build();

// TODO: put databaseName in metadata
get(databaseStub.dropDatabaseCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
get(databaseStub.dropDatabaseCallable().futureCall(request, context));
}

@Override
Expand All @@ -293,25 +308,29 @@ public Database getDatabase(String databaseName) throws SpannerException {
.setName(databaseName)
.build();

// TODO: put databaseName in metadata
return get(databaseStub.getDatabaseCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
return get(databaseStub.getDatabaseCallable().futureCall(request, context));
}

@Override
public List<String> getDatabaseDdl(String databaseName) throws SpannerException {
GetDatabaseDdlRequest request =
GetDatabaseDdlRequest.newBuilder().setDatabase(databaseName).build();

// TODO: put databaseName in metadata
return get(databaseStub.getDatabaseDdlCallable().futureCall(request))
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(databaseName, projectName));
return get(databaseStub.getDatabaseDdlCallable().futureCall(request, context))
.getStatementsList();
}

@Override
public Operation getOperation(String name) throws SpannerException {
GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build();
// TODO: put name in metadata
return get(databaseStub.getOperationsStub().getOperationCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withExtraHeaders(metadataProvider.newExtraHeaders(name, projectName));
return get(databaseStub.getOperationsStub().getOperationCallable()
.futureCall(request, context));
}

@Override
Expand Down Expand Up @@ -365,9 +384,10 @@ public PartitionResponse partitionQuery(
@Override
public PartitionResponse partitionRead(
PartitionReadRequest request, @Nullable Map<Option, ?> options) throws SpannerException {
// TODO(pongad): Figure out metadata
// TODO(pongad): Figure out channel affinity
return get(stub.partitionReadCallable().futureCall(request));
GrpcCallContext context = GrpcCallContext.createDefault()
.withChannelAffinity(Option.CHANNEL_HINT.getLong(options).intValue())
.withExtraHeaders(metadataProvider.newExtraHeaders(request.getSession(), projectName));
return get(stub.partitionReadCallable().futureCall(request, context));
}

/** Gets the result of an async RPC call, handling any exceptions encountered. */
Expand Down
18 changes: 15 additions & 3 deletions ...-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
import com.google.common.collect.ImmutableMap;
import io.grpc.Metadata;
import io.grpc.Metadata.Key;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
Expand All @@ -27,15 +29,15 @@
*/
class SpannerMetadataProvider {
private final Map<Metadata.Key<String>, String> headers;
private final Key<String> resourceHeaderKey;
private final String resourceHeaderKey;

private static final Pattern[] RESOURCE_TOKEN_PATTERNS = {
Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*/databases/[^/]*)(.*)?"),
Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*)(.*)?")
};

private SpannerMetadataProvider(Map<String, String> headers, String resourceHeaderKey) {
this.resourceHeaderKey = Key.of(resourceHeaderKey, Metadata.ASCII_STRING_MARSHALLER);
this.resourceHeaderKey = resourceHeaderKey;
this.headers = constructHeadersAsMetadata(headers);
}

Expand All @@ -50,11 +52,21 @@ Metadata newMetadata(String resourceTokenTemplate, String defaultResourceToken)
}

metadata.put(
resourceHeaderKey, getResourceHeaderValue(resourceTokenTemplate, defaultResourceToken));
Key.of(resourceHeaderKey, Metadata.ASCII_STRING_MARSHALLER),
getResourceHeaderValue(resourceTokenTemplate, defaultResourceToken));

return metadata;
}

Map<String, List<String>> newExtraHeaders(String resourceTokenTemplate, String defaultResourceToken) {
return ImmutableMap.<String, List<String>>builder()
.put(
resourceHeaderKey,
Arrays.asList(getResourceHeaderValue(resourceTokenTemplate, defaultResourceToken)))
.build();

}

private Map<Metadata.Key<String>, String> constructHeadersAsMetadata(
Map<String, String> headers) {
ImmutableMap.Builder<Metadata.Key<String>, String> headersAsMetadataBuilder =
Expand Down

0 comments on commit 9df854c

Please sign in to comment.