From 8594276bae2d884ca69e23e9d2c2c9ac6ed19d1b Mon Sep 17 00:00:00 2001 From: Sakthivel Subramanian Date: Wed, 28 Oct 2026 18:54:34 +0000 Subject: [PATCH] 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();