From d5bc95fd0aec17b1cf4e4d7d8bc698e594bd0159 Mon Sep 17 00:00:00 2001 From: Adam Seering Date: Wed, 10 Jun 2026 19:03:35 +0000 Subject: [PATCH 1/3] feat: port secure_context testing support to executor proxy --- .../executor/spanner/CloudClientExecutor.java | 84 +++++++++++++++++-- 1 file changed, 75 insertions(+), 9 deletions(-) diff --git a/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java b/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java index 1a2e55c1b433..5de241f73f29 100644 --- a/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java +++ b/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java @@ -737,14 +737,16 @@ public synchronized void bufferMutations(List mutations) throws Spanne } /** Execute a batch of updates in a read-write transaction. */ - public synchronized long[] executeBatchDml(@Nonnull List stmts) - throws SpannerException { + public synchronized long[] executeBatchDml( + @Nonnull List stmts, Options.UpdateOption... options) throws SpannerException { for (int i = 0; i < stmts.size(); i++) { LOGGER.log( Level.INFO, String.format("executeBatchDml [%d]: %s", i + 1, stmts.get(i).toString())); } + List allOptions = new ArrayList<>(java.util.Arrays.asList(options)); + allOptions.add(Options.tag("batch-update-transaction-tag")); return getTransactionForWrite() - .batchUpdate(stmts, Options.tag("batch-update-transaction-tag")); + .batchUpdate(stmts, allOptions.toArray(new Options.UpdateOption[0])); } /** Finish active transaction in given finishMode, then send outcome back to client. */ @@ -2199,6 +2201,44 @@ private Status executeGenerateDbPartitionsRead( } /** Execute action that generates database partitions for the given query. */ + private Options.ReadQueryUpdateTransactionOption buildSecureContextOption( + Map secureContextMap) { + if (secureContextMap != null && !secureContextMap.isEmpty()) { + com.google.spanner.v1.RequestOptions.ClientContext.Builder clientContextBuilder = + com.google.spanner.v1.RequestOptions.ClientContext.newBuilder(); + for (Map.Entry entry : + secureContextMap.entrySet()) { + com.google.protobuf.Value.Builder valueBuilder = com.google.protobuf.Value.newBuilder(); + if (entry.getValue().getValueTypeCase() + == com.google.spanner.executor.v1.Value.ValueTypeCase.IS_NULL + && entry.getValue().getIsNull()) { + valueBuilder.setNullValue(com.google.protobuf.NullValue.NULL_VALUE); + } else if (entry.getValue().getValueTypeCase() + == com.google.spanner.executor.v1.Value.ValueTypeCase.STRING_VALUE) { + valueBuilder.setStringValue(entry.getValue().getStringValue()); + } else { + throw new IllegalArgumentException( + "Unsupported secure parameter value type in executor proxy"); + } + clientContextBuilder.putSecureContext(entry.getKey(), valueBuilder.build()); + } + return Options.clientContext(clientContextBuilder.build()); + } + return null; + } + + @SuppressWarnings("unchecked") + private void addSecureContextOption( + Map secureContextMap, + List optionsList, + Class optionClass) { + Options.ReadQueryUpdateTransactionOption secureContextOption = + buildSecureContextOption(secureContextMap); + if (secureContextOption != null && optionClass.isInstance(secureContextOption)) { + optionsList.add((T) secureContextOption); + } + } + private Status executeGenerateDbPartitionsQuery( GenerateDbPartitionsForQueryAction action, OutcomeSender sender, @@ -2206,6 +2246,9 @@ private Status executeGenerateDbPartitionsQuery( try { BatchReadOnlyTransaction batchTxn = executionContext.getBatchTxn(); Statement.Builder stmt = Statement.newBuilder(action.getQuery().getSql()); + List queryOptions = new ArrayList<>(); + addSecureContextOption( + action.getQuery().getSecureContextMap(), queryOptions, Options.QueryOption.class); for (int i = 0; i < action.getQuery().getParamsCount(); ++i) { stmt.bind(action.getQuery().getParams(i).getName()) .to( @@ -2217,7 +2260,9 @@ private Status executeGenerateDbPartitionsQuery( PartitionOptions.newBuilder() .setPartitionSizeBytes(action.getDesiredBytesPerPartition()) .build(); - List parts = batchTxn.partitionQuery(partitionOptions, stmt.build()); + List parts = + batchTxn.partitionQuery( + partitionOptions, stmt.build(), queryOptions.toArray(new Options.QueryOption[0])); List batchPartitions = new ArrayList<>(); for (Partition part : parts) { batchPartitions.add( @@ -2283,11 +2328,15 @@ private Status executePartitionedUpdate( PartitionedUpdateAction action, DatabaseClient dbClient, OutcomeSender sender) { try { ExecutePartitionedUpdateOptions options = action.getOptions(); + List optionsList = new ArrayList<>(); + optionsList.add(Options.tag(options.getTag())); + optionsList.add(Options.priority(RpcPriority.fromProto(options.getRpcPriority()))); + addSecureContextOption( + action.getUpdate().getSecureContextMap(), optionsList, Options.UpdateOption.class); Long count = dbClient.executePartitionedUpdate( Statement.of(action.getUpdate().getSql()), - Options.tag(options.getTag()), - Options.priority(RpcPriority.fromProto(options.getRpcPriority()))); + optionsList.toArray(new Options.UpdateOption[0])); SpannerActionOutcome outcome = SpannerActionOutcome.newBuilder() .setStatus(toProto(Status.OK)) @@ -2740,7 +2789,11 @@ private Status executeQuery( String.format( "Finish query building, ready to execute %s\n", executionContext.getTransactionSeed())); - ResultSet result = txn.executeQuery(stmt.build(), Options.tag("query-tag")); + List queryOptions = new ArrayList<>(); + queryOptions.add(Options.tag("query-tag")); + addSecureContextOption(action.getSecureContextMap(), queryOptions, Options.QueryOption.class); + ResultSet result = + txn.executeQuery(stmt.build(), queryOptions.toArray(new Options.QueryOption[0])); LOGGER.log( Level.INFO, String.format("Parsing query result %s\n", executionContext.getTransactionSeed())); @@ -2770,10 +2823,14 @@ private Status executeCloudDmlUpdate( update.getParams(i).getType(), update.getParams(i).getValue())); } sender.initForQuery(); + List queryOptions = new ArrayList<>(); + queryOptions.add(Options.tag("dml-transaction-tag")); + addSecureContextOption( + action.getUpdate().getSecureContextMap(), queryOptions, Options.QueryOption.class); ResultSet result = executionContext .getTransactionForWrite() - .executeQuery(stmt.build(), Options.tag("dml-transaction-tag")); + .executeQuery(stmt.build(), queryOptions.toArray(new Options.QueryOption[0])); LOGGER.log( Level.INFO, String.format("Parsing Dml result %s\n", executionContext.getTransactionSeed())); @@ -2804,7 +2861,16 @@ private Status executeCloudBatchDmlUpdates( } queries.add(stmt.build()); } - long[] rowCounts = executionContext.executeBatchDml(queries); + Map secureContextMap = + new java.util.HashMap<>(); + for (int i = 0; i < action.getUpdatesCount(); ++i) { + secureContextMap.putAll(action.getUpdates(i).getSecureContextMap()); + } + List optionsList = new ArrayList<>(); + addSecureContextOption(secureContextMap, optionsList, Options.UpdateOption.class); + long[] rowCounts = + executionContext.executeBatchDml( + queries, optionsList.toArray(new Options.UpdateOption[0])); sender.initForQuery(); for (long rowCount : rowCounts) { sender.appendRowsModifiedInDml(rowCount); From 4bb686985d773dee83987d4a377ce11a222937c1 Mon Sep 17 00:00:00 2001 From: Adam Seering Date: Mon, 6 Jul 2026 20:24:46 +0000 Subject: [PATCH 2/3] fix(spanner-executor): port secure_context to SpannerAction and handle empty serviceKeyFile --- .../executor/spanner/CloudClientExecutor.java | 48 +++++++++++-------- .../cloud/executor/spanner/WorkerProxy.java | 12 +++++ 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java b/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java index 5de241f73f29..2684b6573c21 100644 --- a/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java +++ b/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/CloudClientExecutor.java @@ -744,7 +744,16 @@ public synchronized long[] executeBatchDml( Level.INFO, String.format("executeBatchDml [%d]: %s", i + 1, stmts.get(i).toString())); } List allOptions = new ArrayList<>(java.util.Arrays.asList(options)); - allOptions.add(Options.tag("batch-update-transaction-tag")); + boolean hasTag = false; + for (Options.UpdateOption opt : options) { + if (opt.getClass().getSimpleName().equals("TagOption")) { + hasTag = true; + break; + } + } + if (!hasTag) { + allOptions.add(Options.tag("batch-update-tag")); + } return getTransactionForWrite() .batchUpdate(stmts, allOptions.toArray(new Options.UpdateOption[0])); } @@ -863,7 +872,10 @@ private synchronized Spanner initializeClient(long timeoutSeconds, boolean useMu throws IOException { // Create a cloud spanner client Credentials credentials; - if (WorkerProxy.serviceKeyFile.isEmpty()) { + if (WorkerProxy.serviceKeyFile == null + || WorkerProxy.serviceKeyFile.isEmpty() + || !new File(WorkerProxy.serviceKeyFile).exists() + || new File(WorkerProxy.serviceKeyFile).isDirectory()) { credentials = NoCredentials.getInstance(); } else { credentials = @@ -945,7 +957,10 @@ private synchronized TraceServiceClient getTraceServiceClient() throws IOExcepti } // Create a trace service client Credentials credentials; - if (WorkerProxy.serviceKeyFile.isEmpty()) { + if (WorkerProxy.serviceKeyFile == null + || WorkerProxy.serviceKeyFile.isEmpty() + || !new File(WorkerProxy.serviceKeyFile).exists() + || new File(WorkerProxy.serviceKeyFile).isDirectory()) { credentials = NoCredentials.getInstance(); } else { credentials = @@ -2227,15 +2242,13 @@ private Options.ReadQueryUpdateTransactionOption buildSecureContextOption( return null; } - @SuppressWarnings("unchecked") - private void addSecureContextOption( + private void addSecureContextOption( Map secureContextMap, - List optionsList, - Class optionClass) { + List optionsList) { Options.ReadQueryUpdateTransactionOption secureContextOption = buildSecureContextOption(secureContextMap); - if (secureContextOption != null && optionClass.isInstance(secureContextOption)) { - optionsList.add((T) secureContextOption); + if (secureContextOption != null) { + optionsList.add(secureContextOption); } } @@ -2247,8 +2260,7 @@ private Status executeGenerateDbPartitionsQuery( BatchReadOnlyTransaction batchTxn = executionContext.getBatchTxn(); Statement.Builder stmt = Statement.newBuilder(action.getQuery().getSql()); List queryOptions = new ArrayList<>(); - addSecureContextOption( - action.getQuery().getSecureContextMap(), queryOptions, Options.QueryOption.class); + addSecureContextOption(action.getQuery().getSecureContextMap(), queryOptions); for (int i = 0; i < action.getQuery().getParamsCount(); ++i) { stmt.bind(action.getQuery().getParams(i).getName()) .to( @@ -2331,8 +2343,7 @@ private Status executePartitionedUpdate( List optionsList = new ArrayList<>(); optionsList.add(Options.tag(options.getTag())); optionsList.add(Options.priority(RpcPriority.fromProto(options.getRpcPriority()))); - addSecureContextOption( - action.getUpdate().getSecureContextMap(), optionsList, Options.UpdateOption.class); + addSecureContextOption(action.getUpdate().getSecureContextMap(), optionsList); Long count = dbClient.executePartitionedUpdate( Statement.of(action.getUpdate().getSql()), @@ -2791,7 +2802,7 @@ private Status executeQuery( executionContext.getTransactionSeed())); List queryOptions = new ArrayList<>(); queryOptions.add(Options.tag("query-tag")); - addSecureContextOption(action.getSecureContextMap(), queryOptions, Options.QueryOption.class); + addSecureContextOption(action.getSecureContextMap(), queryOptions); ResultSet result = txn.executeQuery(stmt.build(), queryOptions.toArray(new Options.QueryOption[0])); LOGGER.log( @@ -2825,8 +2836,7 @@ private Status executeCloudDmlUpdate( sender.initForQuery(); List queryOptions = new ArrayList<>(); queryOptions.add(Options.tag("dml-transaction-tag")); - addSecureContextOption( - action.getUpdate().getSecureContextMap(), queryOptions, Options.QueryOption.class); + addSecureContextOption(action.getUpdate().getSecureContextMap(), queryOptions); ResultSet result = executionContext .getTransactionForWrite() @@ -2863,11 +2873,11 @@ private Status executeCloudBatchDmlUpdates( } Map secureContextMap = new java.util.HashMap<>(); - for (int i = 0; i < action.getUpdatesCount(); ++i) { - secureContextMap.putAll(action.getUpdates(i).getSecureContextMap()); + for (QueryAction update : action.getUpdatesList()) { + secureContextMap.putAll(update.getSecureContextMap()); } List optionsList = new ArrayList<>(); - addSecureContextOption(secureContextMap, optionsList, Options.UpdateOption.class); + addSecureContextOption(secureContextMap, optionsList); long[] rowCounts = executionContext.executeBatchDml( queries, optionsList.toArray(new Options.UpdateOption[0])); diff --git a/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/WorkerProxy.java b/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/WorkerProxy.java index f56a824ca660..a0bc70d0fd2e 100644 --- a/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/WorkerProxy.java +++ b/java-spanner/google-cloud-spanner-executor/src/main/java/com/google/cloud/executor/spanner/WorkerProxy.java @@ -84,6 +84,18 @@ public class WorkerProxy { private static final double MIN_RATIO = 0.0, MAX_RATIO = 1.0, TRACE_SAMPLING_RATE = 0.01; public static OpenTelemetrySdk setupOpenTelemetrySdk() throws Exception { + if (serviceKeyFile == null + || serviceKeyFile.isEmpty() + || !new File(serviceKeyFile).exists() + || new File(serviceKeyFile).isDirectory()) { + LOGGER.log( + Level.WARNING, + "serviceKeyFile is empty or invalid; starting without OpenTelemetry trace export: " + + serviceKeyFile); + return OpenTelemetrySdk.builder() + .setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance())) + .build(); + } // Read credentials from the serviceKeyFile. HttpTransportFactory HTTP_TRANSPORT_FACTORY = NetHttpTransport::new; Credentials credentials = From a5b1829af14d9645afadea997d6a5eb542aaedd6 Mon Sep 17 00:00:00 2001 From: Adam Seering Date: Tue, 7 Jul 2026 17:19:30 +0000 Subject: [PATCH 3/3] chore: trigger rerun due to stuck split-units shard 8