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..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 @@ -737,14 +737,25 @@ 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)); + 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, 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. */ @@ -861,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 = @@ -943,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 = @@ -2199,6 +2216,42 @@ 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; + } + + private void addSecureContextOption( + Map secureContextMap, + List optionsList) { + Options.ReadQueryUpdateTransactionOption secureContextOption = + buildSecureContextOption(secureContextMap); + if (secureContextOption != null) { + optionsList.add(secureContextOption); + } + } + private Status executeGenerateDbPartitionsQuery( GenerateDbPartitionsForQueryAction action, OutcomeSender sender, @@ -2206,6 +2259,8 @@ 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); for (int i = 0; i < action.getQuery().getParamsCount(); ++i) { stmt.bind(action.getQuery().getParams(i).getName()) .to( @@ -2217,7 +2272,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 +2340,14 @@ 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); 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 +2800,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); + 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 +2834,13 @@ 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); 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 +2871,16 @@ private Status executeCloudBatchDmlUpdates( } queries.add(stmt.build()); } - long[] rowCounts = executionContext.executeBatchDml(queries); + Map secureContextMap = + new java.util.HashMap<>(); + for (QueryAction update : action.getUpdatesList()) { + secureContextMap.putAll(update.getSecureContextMap()); + } + List optionsList = new ArrayList<>(); + addSecureContextOption(secureContextMap, optionsList); + long[] rowCounts = + executionContext.executeBatchDml( + queries, optionsList.toArray(new Options.UpdateOption[0])); sender.initForQuery(); for (long rowCount : rowCounts) { sender.appendRowsModifiedInDml(rowCount); 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 =