Skip to content

Commit 0f81bf0

Browse files
aseeringaseering
authored
feat: port secure_context testing support to executor proxy (#13522)
Add support for our executor proxy, so that tools can drive tests/etc using secure parameters
1 parent b7a8504 commit 0f81bf0

2 files changed

Lines changed: 99 additions & 11 deletions

File tree

‎java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java‎

Lines changed: 87 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -737,14 +737,25 @@ public synchronized void bufferMutations(List<Mutation> mutations) throws Spanne
737737
}
738738

739739
/** Execute a batch of updates in a read-write transaction. */
740-
public synchronized long[] executeBatchDml(@Nonnull List<Statement> stmts)
741-
throws SpannerException {
740+
public synchronized long[] executeBatchDml(
741+
@Nonnull List<Statement> stmts, Options.UpdateOption... options) throws SpannerException {
742742
for (int i = 0; i < stmts.size(); i++) {
743743
LOGGER.log(
744744
Level.INFO, String.format("executeBatchDml [%d]: %s", i + 1, stmts.get(i).toString()));
745745
}
746+
List<Options.UpdateOption> allOptions = new ArrayList<>(java.util.Arrays.asList(options));
747+
boolean hasTag = false;
748+
for (Options.UpdateOption opt : options) {
749+
if (opt.getClass().getSimpleName().equals("TagOption")) {
750+
hasTag = true;
751+
break;
752+
}
753+
}
754+
if (!hasTag) {
755+
allOptions.add(Options.tag("batch-update-tag"));
756+
}
746757
return getTransactionForWrite()
747-
.batchUpdate(stmts, Options.tag("batch-update-transaction-tag"));
758+
.batchUpdate(stmts, allOptions.toArray(new Options.UpdateOption[0]));
748759
}
749760

750761
/** Finish active transaction in given finishMode, then send outcome back to client. */
@@ -861,7 +872,10 @@ private synchronized Spanner initializeClient(long timeoutSeconds, boolean useMu
861872
throws IOException {
862873
// Create a cloud spanner client
863874
Credentials credentials;
864-
if (WorkerProxy.serviceKeyFile.isEmpty()) {
875+
if (WorkerProxy.serviceKeyFile == null
876+
|| WorkerProxy.serviceKeyFile.isEmpty()
877+
|| !new File(WorkerProxy.serviceKeyFile).exists()
878+
|| new File(WorkerProxy.serviceKeyFile).isDirectory()) {
865879
credentials = NoCredentials.getInstance();
866880
} else {
867881
credentials =
@@ -943,7 +957,10 @@ private synchronized TraceServiceClient getTraceServiceClient() throws IOExcepti
943957
}
944958
// Create a trace service client
945959
Credentials credentials;
946-
if (WorkerProxy.serviceKeyFile.isEmpty()) {
960+
if (WorkerProxy.serviceKeyFile == null
961+
|| WorkerProxy.serviceKeyFile.isEmpty()
962+
|| !new File(WorkerProxy.serviceKeyFile).exists()
963+
|| new File(WorkerProxy.serviceKeyFile).isDirectory()) {
947964
credentials = NoCredentials.getInstance();
948965
} else {
949966
credentials =
@@ -2199,13 +2216,51 @@ private Status executeGenerateDbPartitionsRead(
21992216
}
22002217

22012218
/** Execute action that generates database partitions for the given query. */
2219+
private Options.ReadQueryUpdateTransactionOption buildSecureContextOption(
2220+
Map<String, com.google.spanner.executor.v1.Value> secureContextMap) {
2221+
if (secureContextMap != null && !secureContextMap.isEmpty()) {
2222+
com.google.spanner.v1.RequestOptions.ClientContext.Builder clientContextBuilder =
2223+
com.google.spanner.v1.RequestOptions.ClientContext.newBuilder();
2224+
for (Map.Entry<String, com.google.spanner.executor.v1.Value> entry :
2225+
secureContextMap.entrySet()) {
2226+
com.google.protobuf.Value.Builder valueBuilder = com.google.protobuf.Value.newBuilder();
2227+
if (entry.getValue().getValueTypeCase()
2228+
== com.google.spanner.executor.v1.Value.ValueTypeCase.IS_NULL
2229+
&& entry.getValue().getIsNull()) {
2230+
valueBuilder.setNullValue(com.google.protobuf.NullValue.NULL_VALUE);
2231+
} else if (entry.getValue().getValueTypeCase()
2232+
== com.google.spanner.executor.v1.Value.ValueTypeCase.STRING_VALUE) {
2233+
valueBuilder.setStringValue(entry.getValue().getStringValue());
2234+
} else {
2235+
throw new IllegalArgumentException(
2236+
"Unsupported secure parameter value type in executor proxy");
2237+
}
2238+
clientContextBuilder.putSecureContext(entry.getKey(), valueBuilder.build());
2239+
}
2240+
return Options.clientContext(clientContextBuilder.build());
2241+
}
2242+
return null;
2243+
}
2244+
2245+
private void addSecureContextOption(
2246+
Map<String, com.google.spanner.executor.v1.Value> secureContextMap,
2247+
List<? super Options.ReadQueryUpdateTransactionOption> optionsList) {
2248+
Options.ReadQueryUpdateTransactionOption secureContextOption =
2249+
buildSecureContextOption(secureContextMap);
2250+
if (secureContextOption != null) {
2251+
optionsList.add(secureContextOption);
2252+
}
2253+
}
2254+
22022255
private Status executeGenerateDbPartitionsQuery(
22032256
GenerateDbPartitionsForQueryAction action,
22042257
OutcomeSender sender,
22052258
ExecutionFlowContext executionContext) {
22062259
try {
22072260
BatchReadOnlyTransaction batchTxn = executionContext.getBatchTxn();
22082261
Statement.Builder stmt = Statement.newBuilder(action.getQuery().getSql());
2262+
List<Options.QueryOption> queryOptions = new ArrayList<>();
2263+
addSecureContextOption(action.getQuery().getSecureContextMap(), queryOptions);
22092264
for (int i = 0; i < action.getQuery().getParamsCount(); ++i) {
22102265
stmt.bind(action.getQuery().getParams(i).getName())
22112266
.to(
@@ -2217,7 +2272,9 @@ private Status executeGenerateDbPartitionsQuery(
22172272
PartitionOptions.newBuilder()
22182273
.setPartitionSizeBytes(action.getDesiredBytesPerPartition())
22192274
.build();
2220-
List<Partition> parts = batchTxn.partitionQuery(partitionOptions, stmt.build());
2275+
List<Partition> parts =
2276+
batchTxn.partitionQuery(
2277+
partitionOptions, stmt.build(), queryOptions.toArray(new Options.QueryOption[0]));
22212278
List<BatchPartition> batchPartitions = new ArrayList<>();
22222279
for (Partition part : parts) {
22232280
batchPartitions.add(
@@ -2283,11 +2340,14 @@ private Status executePartitionedUpdate(
22832340
PartitionedUpdateAction action, DatabaseClient dbClient, OutcomeSender sender) {
22842341
try {
22852342
ExecutePartitionedUpdateOptions options = action.getOptions();
2343+
List<Options.UpdateOption> optionsList = new ArrayList<>();
2344+
optionsList.add(Options.tag(options.getTag()));
2345+
optionsList.add(Options.priority(RpcPriority.fromProto(options.getRpcPriority())));
2346+
addSecureContextOption(action.getUpdate().getSecureContextMap(), optionsList);
22862347
Long count =
22872348
dbClient.executePartitionedUpdate(
22882349
Statement.of(action.getUpdate().getSql()),
2289-
Options.tag(options.getTag()),
2290-
Options.priority(RpcPriority.fromProto(options.getRpcPriority())));
2350+
optionsList.toArray(new Options.UpdateOption[0]));
22912351
SpannerActionOutcome outcome =
22922352
SpannerActionOutcome.newBuilder()
22932353
.setStatus(toProto(Status.OK))
@@ -2740,7 +2800,11 @@ private Status executeQuery(
27402800
String.format(
27412801
"Finish query building, ready to execute %s\n",
27422802
executionContext.getTransactionSeed()));
2743-
ResultSet result = txn.executeQuery(stmt.build(), Options.tag("query-tag"));
2803+
List<Options.QueryOption> queryOptions = new ArrayList<>();
2804+
queryOptions.add(Options.tag("query-tag"));
2805+
addSecureContextOption(action.getSecureContextMap(), queryOptions);
2806+
ResultSet result =
2807+
txn.executeQuery(stmt.build(), queryOptions.toArray(new Options.QueryOption[0]));
27442808
LOGGER.log(
27452809
Level.INFO,
27462810
String.format("Parsing query result %s\n", executionContext.getTransactionSeed()));
@@ -2770,10 +2834,13 @@ private Status executeCloudDmlUpdate(
27702834
update.getParams(i).getType(), update.getParams(i).getValue()));
27712835
}
27722836
sender.initForQuery();
2837+
List<Options.QueryOption> queryOptions = new ArrayList<>();
2838+
queryOptions.add(Options.tag("dml-transaction-tag"));
2839+
addSecureContextOption(action.getUpdate().getSecureContextMap(), queryOptions);
27732840
ResultSet result =
27742841
executionContext
27752842
.getTransactionForWrite()
2776-
.executeQuery(stmt.build(), Options.tag("dml-transaction-tag"));
2843+
.executeQuery(stmt.build(), queryOptions.toArray(new Options.QueryOption[0]));
27772844
LOGGER.log(
27782845
Level.INFO,
27792846
String.format("Parsing Dml result %s\n", executionContext.getTransactionSeed()));
@@ -2804,7 +2871,16 @@ private Status executeCloudBatchDmlUpdates(
28042871
}
28052872
queries.add(stmt.build());
28062873
}
2807-
long[] rowCounts = executionContext.executeBatchDml(queries);
2874+
Map<String, com.google.spanner.executor.v1.Value> secureContextMap =
2875+
new java.util.HashMap<>();
2876+
for (QueryAction update : action.getUpdatesList()) {
2877+
secureContextMap.putAll(update.getSecureContextMap());
2878+
}
2879+
List<Options.UpdateOption> optionsList = new ArrayList<>();
2880+
addSecureContextOption(secureContextMap, optionsList);
2881+
long[] rowCounts =
2882+
executionContext.executeBatchDml(
2883+
queries, optionsList.toArray(new Options.UpdateOption[0]));
28082884
sender.initForQuery();
28092885
for (long rowCount : rowCounts) {
28102886
sender.appendRowsModifiedInDml(rowCount);

‎java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/WorkerProxy.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,18 @@ public class WorkerProxy {
8484
private static final double MIN_RATIO = 0.0, MAX_RATIO = 1.0, TRACE_SAMPLING_RATE = 0.01;
8585

8686
public static OpenTelemetrySdk setupOpenTelemetrySdk() throws Exception {
87+
if (serviceKeyFile == null
88+
|| serviceKeyFile.isEmpty()
89+
|| !new File(serviceKeyFile).exists()
90+
|| new File(serviceKeyFile).isDirectory()) {
91+
LOGGER.log(
92+
Level.WARNING,
93+
"serviceKeyFile is empty or invalid; starting without OpenTelemetry trace export: "
94+
+ serviceKeyFile);
95+
return OpenTelemetrySdk.builder()
96+
.setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
97+
.build();
98+
}
8799
// Read credentials from the serviceKeyFile.
88100
HttpTransportFactory HTTP_TRANSPORT_FACTORY = NetHttpTransport::new;
89101
Credentials credentials =

0 commit comments

Comments
 (0)