From 8594276bae2d884ca69e23e9d2c2c9ac6ed19d1b Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Wed, 28 Oct 2026 18:54:34 +0000 Subject: [PATCH 1/2] fix(spanner): unflake ITSqlMusicScriptTest KeyAwareChannel incorrectly attached routing hints to ExecuteSqlRequest calls inside Partitioned DML transactions when pinned to an affinity endpoint. The fix tracks Partitioned DML transaction IDs from BeginTransaction options so routeFromRequest(ExecuteSqlRequest) skips attaching routing hints for Partitioned DML queries. --- .../cloud/spanner/spi/v1/KeyAwareChannel.java | 47 ++++++++++++++++- .../spanner/spi/v1/KeyAwareChannelTest.java | 51 +++++++++++++++++++ 2 files changed, 96 insertions(+), 2 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java index 0344188c2ff0..bd0f9edda5c3 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java @@ -121,6 +121,13 @@ final class KeyAwareChannel extends ManagedChannel { // Bound and age out entries in case application code abandons a transaction // without sending Commit/Rollback or otherwise clearing affinity. private final Cache transactionAffinities; + // Tracks the IDs of Partitioned DML transactions. Requests for these transactions must never + // carry a routing hint, as Spanner would then try to execute them on the fastpath, which only + // supports reads and queries and rejects Partitioned DML transactions with + // "PartitionedDml Transactions may only be used to execute DML statements". + // Partitioned DML transactions are never committed or rolled back, so entries age out in the + // same way as abandoned read-write transaction affinities. + private final Cache partitionedDmlTransactions; private final EndpointOverloadCooldownTracker endpointOverloadCooldowns; private KeyAwareChannel( @@ -158,7 +165,8 @@ private KeyAwareChannel( this.lifecycleManager = (endpointCacheFactory == null) ? new EndpointLifecycleManager(endpointCache) : null; this.endpointOverloadCooldowns = endpointOverloadCooldowns; - this.transactionAffinities = newTransactionAffinities(transactionAffinityTicker); + this.transactionAffinities = newTransactionCache(transactionAffinityTicker); + this.partitionedDmlTransactions = newTransactionCache(transactionAffinityTicker); } static KeyAwareChannel create( @@ -201,7 +209,7 @@ static KeyAwareChannel create( transactionAffinityTicker); } - private static Cache newTransactionAffinities(Ticker ticker) { + private static Cache newTransactionCache(Ticker ticker) { return CacheBuilder.newBuilder() .maximumSize(MAX_TRACKED_TRANSACTION_AFFINITIES) .expireAfterAccess(TRANSACTION_AFFINITY_TTL_MINUTES, TimeUnit.MINUTES) @@ -393,6 +401,21 @@ private void clearAffinity(ByteString transactionId) { return; } transactionAffinities.invalidate(transactionId); + partitionedDmlTransactions.invalidate(transactionId); + } + + private void recordPartitionedDmlTransaction(ByteString transactionId) { + if (transactionId == null || transactionId.isEmpty()) { + return; + } + partitionedDmlTransactions.put(transactionId, Boolean.TRUE); + } + + private boolean isPartitionedDmlTransaction(ByteString transactionId) { + if (transactionId == null || transactionId.isEmpty()) { + return false; + } + return partitionedDmlTransactions.getIfPresent(transactionId) != null; } private void recordEndpointCooldown( @@ -737,6 +760,7 @@ static final class KeyAwareClientCall @Nullable private io.grpc.Status cancelledStatus; @Nullable private Metadata cancelledTrailers; private boolean shouldRecordTransactionAffinity; + private boolean beginsPartitionedDmlTransaction; private boolean genericRequest; private boolean reroutedToDefault; private boolean responseMessageReceived; @@ -839,6 +863,8 @@ public void sendMessage(RequestT message) { shouldRecordTransactionAffinity = true; allowDefaultAffinity = true; } + beginsPartitionedDmlTransaction = + reqBuilder.hasOptions() && reqBuilder.getOptions().hasPartitionedDml(); message = (RequestT) reqBuilder.build(); } else if (message instanceof CommitRequest) { CommitRequest request = (CommitRequest) message; @@ -1128,6 +1154,12 @@ void maybeRecordAffinity(ByteString transactionId) { parentChannel.recordAffinity(transactionId, selectedEndpoint, allowDefaultAffinity); } + void maybeRecordPartitionedDmlTransaction(ByteString transactionId) { + if (beginsPartitionedDmlTransaction) { + parentChannel.recordPartitionedDmlTransaction(transactionId); + } + } + void maybeClearAffinity() { parentChannel.clearAffinity(transactionIdToClear); } @@ -1184,6 +1216,14 @@ private RoutingDecision routeFromRequest(ExecuteSqlRequest.Builder reqBuilder) { finder = parentChannel.getOrCreateChannelFinder(databaseId); } boolean preferLeader = preferLeader(reqBuilder.getTransaction()); + if (parentChannel.isPartitionedDmlTransaction(transactionId)) { + // Never attach a routing hint to a Partitioned DML request. A routing hint makes Spanner + // try to execute the statement on the fastpath, which does not support Partitioned DML + // and fails with "PartitionedDml Transactions may only be used to execute DML + // statements". Route the request to the transaction's affinity endpoint (or the default + // endpoint) without a routing hint instead. + return new RoutingDecision(finder, endpoint, databaseId, 0L, preferLeader); + } if (databaseId != null && endpoint == null) { // Transaction IDs only pin routing when a read-write affinity entry exists. Otherwise route // by this request's routing hint without leader bias, as read-only transaction IDs can run @@ -1306,6 +1346,9 @@ public void onMessage(ResponseT message) { if (transactionId != null && call.shouldRecordTransactionAffinity) { call.maybeRecordAffinity(transactionId); } + if (transactionId != null) { + call.maybeRecordPartitionedDmlTransaction(transactionId); + } super.onMessage(message); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java index 21d6ac10dd9a..ec86e46da56a 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java @@ -1386,6 +1386,57 @@ public void readWriteTransactionQueryAttachesRoutingHintWhilePinnedToAffinityEnd assertThat(routingHint.getTabletUid()).isEqualTo(2L); } + @Test + public void partitionedDmlQueryDoesNotAttachRoutingHintWhilePinnedToAffinityEndpoint() + throws Exception { + TestHarness harness = createHarness(); + seedCache(harness, createTwoRangeCacheUpdate()); + ByteString transactionId = ByteString.copyFromUtf8("pdml-no-routing-hint"); + String sql = "DELETE FROM T WHERE K = @k"; + + // Begin a Partitioned DML transaction. This records a transaction affinity in the same way as + // a read-write transaction. + ClientCall beginCall = + harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT); + beginCall.start(new CapturingListener(), new Metadata()); + beginCall.sendMessage( + BeginTransactionRequest.newBuilder() + .setSession(SESSION) + .setOptions( + TransactionOptions.newBuilder() + .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance())) + .build()); + @SuppressWarnings("unchecked") + RecordingClientCall beginDelegate = + (RecordingClientCall) + harness.defaultManagedChannel.latestCall(); + beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build()); + beginDelegate.emitOnClose(Status.OK, new Metadata()); + + // Execute the Partitioned DML statement. Spanner rejects ExecuteSql requests for a Partitioned + // DML transaction that carry a routing hint, because the server fastpath treats them as reads: + // "PartitionedDml Transactions may only be used to execute DML statements". + ClientCall pdmlCall = + harness.channel.newCall(SpannerGrpc.getExecuteStreamingSqlMethod(), CallOptions.DEFAULT); + pdmlCall.start(new CapturingListener(), new Metadata()); + pdmlCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setSql(sql) + .setTransaction(TransactionSelector.newBuilder().setId(transactionId)) + .build()); + + @SuppressWarnings("unchecked") + RecordingClientCall pdmlDelegate = + (RecordingClientCall) + harness.defaultManagedChannel.latestCall(); + assertNotNull(pdmlDelegate.lastMessage); + assertThat(pdmlDelegate.lastMessage.getSql()).isEqualTo(sql); + assertThat(pdmlDelegate.lastMessage.hasRoutingHint()).isFalse(); + assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(0); + assertThat(harness.endpointCache.callCountForAddress("server-b:1234")).isEqualTo(0); + } + @Test public void genericUnimplementedReroutesOnceAndLearnsDefaultOnlyMethod() throws Exception { TestHarness harness = createHarness(); From c8b3743c236aafbd71e6446c08ab0f37975682f8 Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Mon, 2 Nov 2026 14:54:40 +0000 Subject: [PATCH 2/2] Addressed comments --- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 19 ++- .../cloud/spanner/spi/v1/KeyAwareChannel.java | 58 ++----- .../spanner/spi/v1/GapicSpannerRpcTest.java | 155 ++++++++++++++++++ .../spanner/spi/v1/KeyAwareChannelTest.java | 125 ++++++++------ 4 files changed, 258 insertions(+), 99 deletions(-) diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index bb4b78d8e016..473abc904788 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/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)) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java index bd0f9edda5c3..91477342827e 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java @@ -121,13 +121,6 @@ final class KeyAwareChannel extends ManagedChannel { // Bound and age out entries in case application code abandons a transaction // without sending Commit/Rollback or otherwise clearing affinity. private final Cache transactionAffinities; - // Tracks the IDs of Partitioned DML transactions. Requests for these transactions must never - // carry a routing hint, as Spanner would then try to execute them on the fastpath, which only - // supports reads and queries and rejects Partitioned DML transactions with - // "PartitionedDml Transactions may only be used to execute DML statements". - // Partitioned DML transactions are never committed or rolled back, so entries age out in the - // same way as abandoned read-write transaction affinities. - private final Cache partitionedDmlTransactions; private final EndpointOverloadCooldownTracker endpointOverloadCooldowns; private KeyAwareChannel( @@ -165,8 +158,7 @@ private KeyAwareChannel( this.lifecycleManager = (endpointCacheFactory == null) ? new EndpointLifecycleManager(endpointCache) : null; this.endpointOverloadCooldowns = endpointOverloadCooldowns; - this.transactionAffinities = newTransactionCache(transactionAffinityTicker); - this.partitionedDmlTransactions = newTransactionCache(transactionAffinityTicker); + this.transactionAffinities = newTransactionAffinities(transactionAffinityTicker); } static KeyAwareChannel create( @@ -209,7 +201,7 @@ static KeyAwareChannel create( transactionAffinityTicker); } - private static Cache newTransactionCache(Ticker ticker) { + private static Cache newTransactionAffinities(Ticker ticker) { return CacheBuilder.newBuilder() .maximumSize(MAX_TRACKED_TRANSACTION_AFFINITIES) .expireAfterAccess(TRANSACTION_AFFINITY_TTL_MINUTES, TimeUnit.MINUTES) @@ -335,6 +327,10 @@ public String authority() { return authority; } + ManagedChannel defaultChannel() { + return defaultChannel; + } + @Override public ClientCall newCall( MethodDescriptor methodDescriptor, CallOptions callOptions) { @@ -401,21 +397,6 @@ private void clearAffinity(ByteString transactionId) { return; } transactionAffinities.invalidate(transactionId); - partitionedDmlTransactions.invalidate(transactionId); - } - - private void recordPartitionedDmlTransaction(ByteString transactionId) { - if (transactionId == null || transactionId.isEmpty()) { - return; - } - partitionedDmlTransactions.put(transactionId, Boolean.TRUE); - } - - private boolean isPartitionedDmlTransaction(ByteString transactionId) { - if (transactionId == null || transactionId.isEmpty()) { - return false; - } - return partitionedDmlTransactions.getIfPresent(transactionId) != null; } private void recordEndpointCooldown( @@ -760,7 +741,6 @@ static final class KeyAwareClientCall @Nullable private io.grpc.Status cancelledStatus; @Nullable private Metadata cancelledTrailers; private boolean shouldRecordTransactionAffinity; - private boolean beginsPartitionedDmlTransaction; private boolean genericRequest; private boolean reroutedToDefault; private boolean responseMessageReceived; @@ -857,14 +837,14 @@ 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; allowDefaultAffinity = true; } - beginsPartitionedDmlTransaction = - reqBuilder.hasOptions() && reqBuilder.getOptions().hasPartitionedDml(); message = (RequestT) reqBuilder.build(); } else if (message instanceof CommitRequest) { CommitRequest request = (CommitRequest) message; @@ -1154,19 +1134,14 @@ void maybeRecordAffinity(ByteString transactionId) { parentChannel.recordAffinity(transactionId, selectedEndpoint, allowDefaultAffinity); } - void maybeRecordPartitionedDmlTransaction(ByteString transactionId) { - if (beginsPartitionedDmlTransaction) { - parentChannel.recordPartitionedDmlTransaction(transactionId); - } - } - void maybeClearAffinity() { parentChannel.clearAffinity(transactionIdToClear); } private void maybeTrackReadWriteBegin(TransactionSelector selector) { if (selector.getSelectorCase() == TransactionSelector.SelectorCase.BEGIN - && !selector.getBegin().hasReadOnly()) { + && !selector.getBegin().hasReadOnly() + && !selector.getBegin().hasPartitionedDml()) { shouldRecordTransactionAffinity = true; allowDefaultAffinity = true; } @@ -1216,14 +1191,6 @@ private RoutingDecision routeFromRequest(ExecuteSqlRequest.Builder reqBuilder) { finder = parentChannel.getOrCreateChannelFinder(databaseId); } boolean preferLeader = preferLeader(reqBuilder.getTransaction()); - if (parentChannel.isPartitionedDmlTransaction(transactionId)) { - // Never attach a routing hint to a Partitioned DML request. A routing hint makes Spanner - // try to execute the statement on the fastpath, which does not support Partitioned DML - // and fails with "PartitionedDml Transactions may only be used to execute DML - // statements". Route the request to the transaction's affinity endpoint (or the default - // endpoint) without a routing hint instead. - return new RoutingDecision(finder, endpoint, databaseId, 0L, preferLeader); - } if (databaseId != null && endpoint == null) { // Transaction IDs only pin routing when a read-write affinity entry exists. Otherwise route // by this request's routing hint without leader bias, as read-only transaction IDs can run @@ -1346,9 +1313,6 @@ public void onMessage(ResponseT message) { if (transactionId != null && call.shouldRecordTransactionAffinity) { call.maybeRecordAffinity(transactionId); } - if (transactionId != null) { - call.maybeRecordPartitionedDmlTransaction(transactionId); - } super.onMessage(message); } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java index 52d5833489b8..c2e4bb601665 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java +++ b/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 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); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java index ec86e46da56a..90b9f39e48d0 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java @@ -1386,57 +1386,6 @@ public void readWriteTransactionQueryAttachesRoutingHintWhilePinnedToAffinityEnd assertThat(routingHint.getTabletUid()).isEqualTo(2L); } - @Test - public void partitionedDmlQueryDoesNotAttachRoutingHintWhilePinnedToAffinityEndpoint() - throws Exception { - TestHarness harness = createHarness(); - seedCache(harness, createTwoRangeCacheUpdate()); - ByteString transactionId = ByteString.copyFromUtf8("pdml-no-routing-hint"); - String sql = "DELETE FROM T WHERE K = @k"; - - // Begin a Partitioned DML transaction. This records a transaction affinity in the same way as - // a read-write transaction. - ClientCall beginCall = - harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT); - beginCall.start(new CapturingListener(), new Metadata()); - beginCall.sendMessage( - BeginTransactionRequest.newBuilder() - .setSession(SESSION) - .setOptions( - TransactionOptions.newBuilder() - .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance())) - .build()); - @SuppressWarnings("unchecked") - RecordingClientCall beginDelegate = - (RecordingClientCall) - harness.defaultManagedChannel.latestCall(); - beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build()); - beginDelegate.emitOnClose(Status.OK, new Metadata()); - - // Execute the Partitioned DML statement. Spanner rejects ExecuteSql requests for a Partitioned - // DML transaction that carry a routing hint, because the server fastpath treats them as reads: - // "PartitionedDml Transactions may only be used to execute DML statements". - ClientCall pdmlCall = - harness.channel.newCall(SpannerGrpc.getExecuteStreamingSqlMethod(), CallOptions.DEFAULT); - pdmlCall.start(new CapturingListener(), new Metadata()); - pdmlCall.sendMessage( - ExecuteSqlRequest.newBuilder() - .setSession(SESSION) - .setSql(sql) - .setTransaction(TransactionSelector.newBuilder().setId(transactionId)) - .build()); - - @SuppressWarnings("unchecked") - RecordingClientCall pdmlDelegate = - (RecordingClientCall) - harness.defaultManagedChannel.latestCall(); - assertNotNull(pdmlDelegate.lastMessage); - assertThat(pdmlDelegate.lastMessage.getSql()).isEqualTo(sql); - assertThat(pdmlDelegate.lastMessage.hasRoutingHint()).isFalse(); - assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(0); - assertThat(harness.endpointCache.callCountForAddress("server-b:1234")).isEqualTo(0); - } - @Test public void genericUnimplementedReroutesOnceAndLearnsDefaultOnlyMethod() throws Exception { TestHarness harness = createHarness(); @@ -2023,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 beginCall = + harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT); + beginCall.start(new CapturingListener(), new Metadata()); + beginCall.sendMessage( + BeginTransactionRequest.newBuilder() + .setSession(SESSION) + .setOptions( + TransactionOptions.newBuilder() + .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance())) + .build()); + + @SuppressWarnings("unchecked") + RecordingClientCall beginDelegate = + (RecordingClientCall) + 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 sqlCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + sqlCall.start(new CapturingListener(), 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 inlineBeginCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + inlineBeginCall.start(new CapturingListener(), new Metadata()); + inlineBeginCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setTransaction( + TransactionSelector.newBuilder() + .setBegin( + TransactionOptions.newBuilder() + .setPartitionedDml( + TransactionOptions.PartitionedDml.getDefaultInstance()))) + .build()); + @SuppressWarnings("unchecked") + RecordingClientCall inlineBeginDelegate = + (RecordingClientCall) + harness.defaultManagedChannel.latestCall(); + inlineBeginDelegate.emitOnMessage( + ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setTransaction(Transaction.newBuilder().setId(inlinePdmlId))) + .build()); + inlineBeginDelegate.emitOnClose(Status.OK, new Metadata()); + + ClientCall followUpCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + followUpCall.start(new CapturingListener(), 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();