feat: port secure_context testing support to executor proxy

Author: aseering

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

@@ -2198,13 +2198,51 @@ private Status executeGenerateDbPartitionsRead(
   }
 
   /** Execute action that generates database partitions for the given query. */
+
+  private Options.ReadQueryUpdateTransactionOption buildSecureContextOption(
+      Map<String, com.google.spanner.executor.v1.Value> secureContextMap) {
+    if (secureContextMap != null && !secureContextMap.isEmpty()) {
+      com.google.spanner.v1.RequestOptions.ClientContext.Builder clientContextBuilder =
+          com.google.spanner.v1.RequestOptions.ClientContext.newBuilder();
+      for (Map.Entry<String, com.google.spanner.executor.v1.Value> 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 GitHub proxy");
+        }
+        clientContextBuilder.putSecureContext(entry.getKey(), valueBuilder.build());
+      }
+      return Options.clientContext(clientContextBuilder.build());
+    }
+    return null;
+  }
+
+  @SuppressWarnings("unchecked")
+  private <T> void addSecureContextOption(
+      Map<String, com.google.spanner.executor.v1.Value> secureContextMap,
+      List<T> 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,
       ExecutionFlowContext executionContext) {
     try {
       BatchReadOnlyTransaction batchTxn = executionContext.getBatchTxn();
       Statement.Builder stmt = Statement.newBuilder(action.getQuery().getSql());
+      List<Options.QueryOption> 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(
@@ -2216,7 +2254,7 @@ private Status executeGenerateDbPartitionsQuery(
           PartitionOptions.newBuilder()
               .setPartitionSizeBytes(action.getDesiredBytesPerPartition())
               .build();
-      List<Partition> parts = batchTxn.partitionQuery(partitionOptions, stmt.build());
+      List<Partition> parts = batchTxn.partitionQuery(partitionOptions, stmt.build(), queryOptions.toArray(new Options.QueryOption[0]));
       List<BatchPartition> batchPartitions = new ArrayList<>();
       for (Partition part : parts) {
         batchPartitions.add(
@@ -2282,11 +2320,14 @@ private Status executePartitionedUpdate(
       PartitionedUpdateAction action, DatabaseClient dbClient, OutcomeSender sender) {
     try {
       ExecutePartitionedUpdateOptions options = action.getOptions();
+      List<Options.UpdateOption> 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))
@@ -2739,7 +2780,10 @@ 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<Options.QueryOption> 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()));

java-spanner/proto-google-cloud-spanner-executor-v1/src/main/java/com/google/spanner/executor/v1/QueryAction.java

@@ -1380,7 +1380,8 @@ public com.google.protobuf.ByteString getSqlBytes() {
    * <code>repeated .google.spanner.executor.v1.QueryAction.Parameter params = 2;</code>
    */
   @java.lang.Override
-  public java.util.List<com.google.spanner.executor.v1.QueryAction.Parameter> getParamsList() {
+  public java.util.Map<java.lang.String, com.google.spanner.executor.v1.Value> getSecureContextMap() { return java.util.Collections.emptyMap(); }
+    public java.util.List<com.google.spanner.executor.v1.QueryAction.Parameter> getParamsList() {
     return params_;
   }
 
@@ -1960,6 +1961,7 @@ private void ensureParamsIsMutable() {
      *
      * <code>repeated .google.spanner.executor.v1.QueryAction.Parameter params = 2;</code>
      */
+    public java.util.Map<java.lang.String, com.google.spanner.executor.v1.Value> getSecureContextMap() { return java.util.Collections.emptyMap(); }
     public java.util.List<com.google.spanner.executor.v1.QueryAction.Parameter> getParamsList() {
       if (paramsBuilder_ == null) {
         return java.util.Collections.unmodifiableList(params_);

java-spanner/proto-google-cloud-spanner-executor-v1/src/main/java/com/google/spanner/executor/v1/QueryActionOrBuilder.java

@@ -61,6 +61,7 @@ public interface QueryActionOrBuilder
    *
    * <code>repeated .google.spanner.executor.v1.QueryAction.Parameter params = 2;</code>
    */
+  java.util.Map<java.lang.String, com.google.spanner.executor.v1.Value> getSecureContextMap();
   java.util.List<com.google.spanner.executor.v1.QueryAction.Parameter> getParamsList();
 
   /**

java-spanner/proto-google-cloud-spanner-executor-v1/src/main/proto/google/spanner/executor/v1/cloud_executor.proto

@@ -178,6 +178,8 @@ message QueryAction {
 
   // Parameters for the SQL string.
   repeated Parameter params = 2;
+n  // Secure context parameters.
+  map<string, Value> secure_context = 3;
 }
 
 // A single DML statement.