test(spanner): unflake ITSqlMusicScriptTest

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -534,9 +534,11 @@ public GapicSpannerRpc(final SpannerOptions options) {
                   .getStreamWatchdogProvider()
                   .withCheckInterval(pdmlSettings.getStreamWatchdogCheckInterval()));
         }
+        ClientContext partitionedDmlClientContext =
+            createPartitionedDmlClientContext(clientContext, this.keyAwareChannel);
         this.partitionedDmlStub =
             GrpcSpannerStubWithStubSettingsAndClientContext.create(
-                pdmlSettings.build(), clientContext);
+                pdmlSettings.build(), partitionedDmlClientContext);
         this.instanceAdminStubSettings =
             options.getInstanceAdminStubSettings().toBuilder()
                 .setTransportChannelProvider(channelProvider)
@@ -677,6 +679,21 @@ private static KeyAwareChannel extractKeyAwareChannel(TransportChannel transport
     return null;
   }
 
+  private static ClientContext createPartitionedDmlClientContext(
+      ClientContext clientContext, @Nullable KeyAwareChannel keyAwareChannel) {
+    if (keyAwareChannel == null) {
+      return clientContext;
+    }
+    GrpcTransportChannel defaultTransportChannel =
+        GrpcTransportChannel.create(keyAwareChannel.defaultChannel());
+    return clientContext.toBuilder()
+        .setTransportChannel(defaultTransportChannel)
+        .setDefaultCallContext(
+            clientContext.getDefaultCallContext().withTransportChannel(defaultTransportChannel))
+        .setBackgroundResources(Collections.emptyList())
+        .build();
+  }
+
   @Nullable
   private static GcpManagedChannel extractGrpcGcpChannel(TransportChannel transportChannel) {
     if (!(transportChannel instanceof GrpcTransportChannel)) {

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java

@@ -327,6 +327,10 @@ public String authority() {
     return authority;
   }
 
+  ManagedChannel defaultChannel() {
+    return defaultChannel;
+  }
+
   @Override
   public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
       MethodDescriptor<RequestT, ResponseT> methodDescriptor, CallOptions callOptions) {
@@ -833,7 +837,9 @@ public void sendMessage(RequestT message) {
           if (finder != null && reqBuilder.hasMutationKey()) {
             endpoint = finder.findServerAndFillRoutingHint(reqBuilder, excludedEndpoints);
           }
-          if (reqBuilder.hasOptions() && reqBuilder.getOptions().hasReadOnly()) {
+          if (reqBuilder.hasOptions()
+              && (reqBuilder.getOptions().hasReadOnly()
+                  || reqBuilder.getOptions().hasPartitionedDml())) {
             shouldRecordTransactionAffinity = false;
           } else {
             shouldRecordTransactionAffinity = true;
@@ -1134,7 +1140,8 @@ void maybeClearAffinity() {
 
     private void maybeTrackReadWriteBegin(TransactionSelector selector) {
       if (selector.getSelectorCase() == TransactionSelector.SelectorCase.BEGIN
-          && !selector.getBegin().hasReadOnly()) {
+          && !selector.getBegin().hasReadOnly()
+          && !selector.getBegin().hasPartitionedDml()) {
         shouldRecordTransactionAffinity = true;
         allowDefaultAffinity = true;
       }

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java

@@ -75,24 +75,30 @@
 import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.Futures;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.ListValue;
 import com.google.rpc.ErrorInfo;
 import com.google.spanner.admin.database.v1.DatabaseAdminGrpc;
 import com.google.spanner.admin.database.v1.ListDatabaseRolesRequest;
+import com.google.spanner.v1.BeginTransactionRequest;
 import com.google.spanner.v1.CommitRequest;
 import com.google.spanner.v1.CreateSessionRequest;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.GetSessionRequest;
+import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.ReadRequest;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.Session;
 import com.google.spanner.v1.SpannerGrpc;
 import com.google.spanner.v1.StructType;
 import com.google.spanner.v1.StructType.Field;
+import com.google.spanner.v1.TransactionOptions;
+import com.google.spanner.v1.TransactionSelector;
 import com.google.spanner.v1.TypeCode;
 import io.grpc.CallOptions;
 import io.grpc.Context;
 import io.grpc.Contexts;
+import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.Metadata;
 import io.grpc.Metadata.Key;
@@ -141,6 +147,7 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BooleanSupplier;
 import javax.annotation.Nullable;
 import org.junit.After;
@@ -2208,6 +2215,154 @@ public boolean isEnableLocationApi() {
     }
   }
 
+  @Test
+  public void testPartitionedDmlUsesDefaultChannelWithoutEnteringKeyAwareChannel()
+      throws Exception {
+    AtomicInteger keyAwareRoutedCalls = new AtomicInteger();
+    ChannelEndpointCacheFactory factory =
+        baseProvider -> {
+          GrpcChannelEndpointCache delegateCache = new GrpcChannelEndpointCache(baseProvider);
+          ChannelEndpoint delegateDefault = delegateCache.defaultChannel();
+          ChannelEndpoint countingDefault =
+              new ChannelEndpoint() {
+                @Override
+                public String getAddress() {
+                  return delegateDefault.getAddress();
+                }
+
+                @Override
+                public boolean isHealthy() {
+                  return delegateDefault.isHealthy();
+                }
+
+                @Override
+                public boolean isTransientFailure() {
+                  return delegateDefault.isTransientFailure();
+                }
+
+                @Override
+                public ManagedChannel getChannel() {
+                  return delegateDefault.getChannel();
+                }
+
+                @Override
+                public void incrementActiveRequests() {
+                  keyAwareRoutedCalls.incrementAndGet();
+                  delegateDefault.incrementActiveRequests();
+                }
+
+                @Override
+                public void decrementActiveRequests() {
+                  delegateDefault.decrementActiveRequests();
+                }
+
+                @Override
+                public int getActiveRequestCount() {
+                  return delegateDefault.getActiveRequestCount();
+                }
+              };
+          return new ChannelEndpointCache() {
+            @Override
+            public ChannelEndpoint defaultChannel() {
+              return countingDefault;
+            }
+
+            @Override
+            public ChannelEndpoint get(String addr) {
+              return addr.equals(countingDefault.getAddress())
+                  ? countingDefault
+                  : delegateCache.get(addr);
+            }
+
+            @Override
+            @Nullable
+            public ChannelEndpoint getIfPresent(String addr) {
+              return addr.equals(countingDefault.getAddress())
+                  ? countingDefault
+                  : delegateCache.getIfPresent(addr);
+            }
+
+            @Override
+            public void evict(String addr) {
+              delegateCache.evict(addr);
+            }
+
+            @Override
+            public void shutdown() {
+              delegateCache.shutdown();
+            }
+          };
+        };
+
+    try {
+      SpannerOptions.useEnvironment(
+          new SpannerOptions.SpannerEnvironment() {
+            @Override
+            public boolean isEnableLocationApi() {
+              return true;
+            }
+          });
+      SpannerOptions options =
+          createSpannerOptions().toBuilder().setChannelEndpointCacheFactory(factory).build();
+      try (Spanner testSpanner = options.getService()) {
+        DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "i", "d"));
+        try (ResultSet rs = client.singleUse().executeQuery(SELECT1AND2)) {
+          while (rs.next()) {}
+        }
+        assertThat(keyAwareRoutedCalls.get()).isGreaterThan(0);
+
+        mockSpanner.clearRequests();
+        keyAwareRoutedCalls.set(0);
+
+        // End-to-end Partitioned DML: BeginTransaction goes through spannerStub (1 call), while
+        // ExecuteStreamingSql goes through partitionedDmlStub directly on the default channel (0
+        // KeyAwareClientCall entries).
+        long updateCount = client.executePartitionedUpdate(UPDATE_FOO_STATEMENT);
+        assertEquals(1L, updateCount);
+        assertEquals(1, keyAwareRoutedCalls.get());
+        List<ExecuteSqlRequest> pdmlStreamingRequests =
+            mockSpanner.getRequestsOfType(ExecuteSqlRequest.class);
+        assertEquals(1, pdmlStreamingRequests.size());
+        assertFalse(pdmlStreamingRequests.get(0).hasRoutingHint());
+
+        // Direct unary and streaming Partitioned DML RPCs must never enter KeyAwareChannel at all.
+        GapicSpannerRpc rpc = getRpc(testSpanner);
+        String sessionName =
+            mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0).getSession();
+        BeginTransactionRequest beginPdmlRequest =
+            BeginTransactionRequest.newBuilder()
+                .setSession(sessionName)
+                .setOptions(
+                    TransactionOptions.newBuilder()
+                        .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance()))
+                .build();
+        ByteString unaryTxId = rpc.beginTransaction(beginPdmlRequest, optionsMap, true).getId();
+        ByteString streamingTxId = rpc.beginTransaction(beginPdmlRequest, optionsMap, true).getId();
+        keyAwareRoutedCalls.set(0);
+        rpc.executePartitionedDml(
+            ExecuteSqlRequest.newBuilder()
+                .setSession(sessionName)
+                .setSql(UPDATE_FOO_STATEMENT.getSql())
+                .setTransaction(TransactionSelector.newBuilder().setId(unaryTxId))
+                .build(),
+            optionsMap);
+        for (PartialResultSet ignored :
+            rpc.executeStreamingPartitionedDml(
+                ExecuteSqlRequest.newBuilder()
+                    .setSession(sessionName)
+                    .setSql(UPDATE_FOO_STATEMENT.getSql())
+                    .setTransaction(TransactionSelector.newBuilder().setId(streamingTxId))
+                    .build(),
+                optionsMap,
+                rpc.getRequestIdCreator().nextRequestId(0),
+                Duration.ofSeconds(10))) {}
+        assertEquals(0, keyAwareRoutedCalls.get());
+      }
+    } finally {
+      SpannerOptions.useDefaultEnvironment();
+    }
+  }
+
   @Test
   public void testLocationApiDoesNotOverrideExplicitChannelProvider() {
     AtomicBoolean factoryCalled = new AtomicBoolean(false);

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java

@@ -1972,6 +1972,80 @@ public void readOnlyTransactionDoesNotRecordAffinity() throws Exception {
     assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);
   }
 
+  @Test
+  public void partitionedDmlTransactionDoesNotRecordAffinity() throws Exception {
+    TestHarness harness = createHarness();
+    ByteString transactionId = ByteString.copyFromUtf8("pdml-tx-no-affinity");
+
+    // Begin a Partitioned DML transaction via BeginTransaction.
+    ClientCall<BeginTransactionRequest, Transaction> beginCall =
+        harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
+    beginCall.start(new CapturingListener<Transaction>(), new Metadata());
+    beginCall.sendMessage(
+        BeginTransactionRequest.newBuilder()
+            .setSession(SESSION)
+            .setOptions(
+                TransactionOptions.newBuilder()
+                    .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance()))
+            .build());
+
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
+        (RecordingClientCall<BeginTransactionRequest, Transaction>)
+            harness.defaultManagedChannel.latestCall();
+    beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
+    beginDelegate.emitOnClose(Status.OK, new Metadata());
+
+    // Send a request referencing the Partitioned DML transaction ID. Because no affinity was
+    // recorded on BeginTransaction, the endpoint cache must not be queried for affinity lookup.
+    ClientCall<ExecuteSqlRequest, ResultSet> sqlCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    sqlCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    sqlCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(transactionId))
+            .build());
+    assertThat(harness.endpointCache.getIfPresentCount(DEFAULT_ADDRESS)).isEqualTo(0);
+
+    // Also verify inline TransactionSelector.begin with PartitionedDml does not record affinity.
+    ByteString inlinePdmlId = ByteString.copyFromUtf8("pdml-inline-no-affinity");
+    ClientCall<ExecuteSqlRequest, ResultSet> inlineBeginCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    inlineBeginCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    inlineBeginCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(
+                TransactionSelector.newBuilder()
+                    .setBegin(
+                        TransactionOptions.newBuilder()
+                            .setPartitionedDml(
+                                TransactionOptions.PartitionedDml.getDefaultInstance())))
+            .build());
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<ExecuteSqlRequest, ResultSet> inlineBeginDelegate =
+        (RecordingClientCall<ExecuteSqlRequest, ResultSet>)
+            harness.defaultManagedChannel.latestCall();
+    inlineBeginDelegate.emitOnMessage(
+        ResultSet.newBuilder()
+            .setMetadata(
+                ResultSetMetadata.newBuilder()
+                    .setTransaction(Transaction.newBuilder().setId(inlinePdmlId)))
+            .build());
+    inlineBeginDelegate.emitOnClose(Status.OK, new Metadata());
+
+    ClientCall<ExecuteSqlRequest, ResultSet> followUpCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    followUpCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    followUpCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(inlinePdmlId))
+            .build());
+    assertThat(harness.endpointCache.getIfPresentCount(DEFAULT_ADDRESS)).isEqualTo(0);
+  }
+
   @Test
   public void abandonedReadWriteTransactionAffinityExpiresAfterInactivity() throws Exception {
     FakeTicker ticker = new FakeTicker();