fix: resolve comments and add new tests to verify that the route-to-leader header exists for RW transactions and does not exist for RO transactions or when the leader aware routing feature is disabled.

yifanzyifanz · yifanzyifanz · commit 1bbd24dc53a2 · 2023-02-01T15:27:30.000-08:00
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java
@@ -161,7 +161,11 @@ static Builder newBuilder() {
     private SingleReadContext(Builder builder) {
       super(builder);
       this.bound = builder.bound;
-      this.routeToLeader = false;
+    }
+
+    @Override
+    protected boolean isRouteToLeader() {
+      return false;
     }
 
     @GuardedBy("lock")
@@ -292,7 +296,11 @@ static Builder newBuilder() {
         this.timestamp = builder.timestamp;
         this.transactionId = builder.transactionId;
       }
-      this.routeToLeader = false;
+    }
+
+    @Override
+    protected boolean isRouteToLeader() {
+      return false;
     }
 
     @Override
@@ -350,7 +358,7 @@ void initTransaction() {
                   .setOptions(options)
                   .build();
           Transaction transaction =
-              rpc.beginTransaction(request, session.getOptions(), routeToLeader);
+              rpc.beginTransaction(request, session.getOptions(), isRouteToLeader());
           if (!transaction.hasReadTimestamp()) {
             throw SpannerExceptionFactory.newSpannerException(
                 ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field");
@@ -383,7 +391,6 @@ void initTransaction() {
   Span span;
   private final int defaultPrefetchChunks;
   private final QueryOptions defaultQueryOptions;
-  protected boolean routeToLeader = false;
 
   @GuardedBy("lock")
   private boolean isValid = true;
@@ -420,6 +427,10 @@ long getSeqNo() {
     return seqNo.incrementAndGet();
   }
 
+  protected boolean isRouteToLeader() {
+    return false;
+  }
+
   @Override
   public final ResultSet read(
       String table, KeySet keys, Iterable<String> columns, ReadOption... options) {
@@ -669,7 +680,7 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
             }
             SpannerRpc.StreamingCall call =
                 rpc.executeQuery(
-                    request.build(), stream.consumer(), session.getOptions(), routeToLeader);
+                    request.build(), stream.consumer(), session.getOptions(), isRouteToLeader());
             call.request(prefetchChunks);
             stream.setCall(call, request.getTransaction().hasBegin());
             return stream;
@@ -797,7 +808,8 @@ CloseableIterator<PartialResultSet> startStream(@Nullable ByteString resumeToken
             }
             builder.setRequestOptions(buildRequestOptions(readOptions));
             SpannerRpc.StreamingCall call =
-                rpc.read(builder.build(), stream.consumer(), session.getOptions(), routeToLeader);
+                rpc.read(
+                    builder.build(), stream.consumer(), session.getOptions(), isRouteToLeader());
             call.request(prefetchChunks);
             stream.setCall(call, /* withBeginTransaction = */ builder.getTransaction().hasBegin());
             return stream;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java
@@ -196,7 +196,11 @@ private TransactionContextImpl(Builder builder) {
       this.trackTransactionStarter = builder.trackTransactionStarter;
       this.options = builder.options;
       this.finishedAsyncOperations.set(null);
-      this.routeToLeader = true;
+    }
+
+    @Override
+    protected boolean isRouteToLeader() {
+      return true;
     }
 
     private void increaseAsyncOperations() {
@@ -256,7 +260,7 @@ ApiFuture<Void> ensureTxnAsync() {
 
     private void createTxnAsync(final SettableApiFuture<Void> res) {
       span.addAnnotation("Creating Transaction");
-      final ApiFuture<ByteString> fut = session.beginTransactionAsync(options, routeToLeader);
+      final ApiFuture<ByteString> fut = session.beginTransactionAsync(options, isRouteToLeader());
       fut.addListener(
           () -> {
             try {
@@ -718,7 +722,7 @@ private ResultSet internalExecuteUpdate(
               /* withTransactionSelector = */ true);
       try {
         com.google.spanner.v1.ResultSet resultSet =
-            rpc.executeQuery(builder.build(), session.getOptions(), routeToLeader);
+            rpc.executeQuery(builder.build(), session.getOptions(), isRouteToLeader());
         if (resultSet.getMetadata().hasTransaction()) {
           onTransactionMetadata(
               resultSet.getMetadata().getTransaction(), builder.getTransaction().hasBegin());
@@ -748,7 +752,7 @@ public ApiFuture<Long> executeUpdateAsync(Statement statement, UpdateOption... o
         // Register the update as an async operation that must finish before the transaction may
         // commit.
         increaseAsyncOperations();
-        resultSet = rpc.executeQueryAsync(builder.build(), session.getOptions(), routeToLeader);
+        resultSet = rpc.executeQueryAsync(builder.build(), session.getOptions(), isRouteToLeader());
       } catch (Throwable t) {
         decreaseAsyncOperations();
         throw t;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerMetadataProvider.java
@@ -28,12 +28,15 @@
 class SpannerMetadataProvider {
   private final Map<Metadata.Key<String>, String> headers;
   private final String resourceHeaderKey;
-  private static final String routeToLeaderHeaderKey = "x-goog-spanner-route-to-leader";
+  private static final String ROUTE_TO_LEADER_HEADER_KEY = "x-goog-spanner-route-to-leader";
   private static final Pattern[] RESOURCE_TOKEN_PATTERNS = {
     Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*/databases/[^/]*)(.*)?"),
     Pattern.compile("^(?<headerValue>projects/[^/]*/instances/[^/]*)(.*)?")
   };
 
+  private static final Map<String, List<String>> ROUTE_TO_LEADER_HEADER_MAP =
+      ImmutableMap.of(ROUTE_TO_LEADER_HEADER_KEY, Collections.singletonList("true"));
+
   private SpannerMetadataProvider(Map<String, String> headers, String resourceHeaderKey) {
     this.resourceHeaderKey = resourceHeaderKey;
     this.headers = constructHeadersAsMetadata(headers);
@@ -67,9 +70,7 @@ Map<String, List<String>> newExtraHeaders(
   }
 
   Map<String, List<String>> newRouteToLeaderHeader() {
-    return ImmutableMap.<String, List<String>>builder()
-        .put(routeToLeaderHeaderKey, Collections.singletonList("true"))
-        .build();
+    return ROUTE_TO_LEADER_HEADER_MAP;
   }
 
   private Map<Metadata.Key<String>, String> constructHeadersAsMetadata(
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java
@@ -19,9 +19,11 @@
 import static com.google.common.truth.Truth.assertThat;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
 import com.google.api.gax.core.GaxProperties;
@@ -45,6 +47,7 @@
 import com.google.cloud.spanner.SpannerOptions;
 import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator;
 import com.google.cloud.spanner.Statement;
+import com.google.cloud.spanner.TransactionRunner;
 import com.google.cloud.spanner.spi.v1.GapicSpannerRpc.AdminRequestsLimitExceededRetryAlgorithm;
 import com.google.cloud.spanner.spi.v1.SpannerRpc.Option;
 import com.google.common.collect.ImmutableList;
@@ -140,6 +143,7 @@ public class GapicSpannerRpcTest {
   private static Metadata lastSeenHeaders;
   private static String defaultUserAgent;
   private static Spanner spanner;
+  private static boolean isRouteToLeader;
 
   @Parameter public Dialect dialect;
 
@@ -177,6 +181,17 @@ public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(
                     String auth =
                         headers.get(Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER));
                     assertThat(auth).isEqualTo("Bearer " + VARIABLE_OAUTH_TOKEN);
+                    if (call.getMethodDescriptor()
+                            .equals(SpannerGrpc.getExecuteStreamingSqlMethod())
+                        || call.getMethodDescriptor().equals(SpannerGrpc.getExecuteSqlMethod())) {
+                      String routeToLeaderHeader =
+                          headers.get(
+                              Key.of(
+                                  "x-goog-spanner-route-to-leader",
+                                  Metadata.ASCII_STRING_MARSHALLER));
+                      isRouteToLeader =
+                          (routeToLeaderHeader != null && routeToLeaderHeader.equals("true"));
+                    }
                     return Contexts.interceptCall(Context.current(), call, headers, next);
                   }
                 })
@@ -198,6 +213,7 @@ public void reset() throws InterruptedException {
       server.shutdown();
       server.awaitTermination();
     }
+    isRouteToLeader = false;
   }
 
   @Test
@@ -508,6 +524,54 @@ public void testCustomUserAgent() {
     }
   }
 
+  @Test
+  public void testRouteToLeaderHeaderForReadOnly() {
+    final SpannerOptions options = createSpannerOptions();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+
+      try (final ResultSet rs = databaseClient.singleUse().executeQuery(SELECT1AND2)) {
+        rs.next();
+      }
+
+      assertFalse(isRouteToLeader);
+    }
+  }
+
+  @Test
+  public void testRouteToLeaderHeaderForReadWrite() {
+    final SpannerOptions options = createSpannerOptions();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+      TransactionRunner runner = databaseClient.readWriteTransaction();
+      runner.run(
+          transaction -> {
+            transaction.executeUpdate(UPDATE_FOO_STATEMENT);
+            return null;
+          });
+    }
+    assertTrue(isRouteToLeader);
+  }
+
+  @Test
+  public void testRouteToLeaderHeaderWithLeaderAwareRoutingDisabled() {
+    final SpannerOptions options =
+        createSpannerOptions().toBuilder().disableLeaderAwareRouting().build();
+    try (Spanner spanner = options.getService()) {
+      final DatabaseClient databaseClient =
+          spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+      TransactionRunner runner = databaseClient.readWriteTransaction();
+      runner.run(
+          transaction -> {
+            transaction.executeUpdate(UPDATE_FOO_STATEMENT);
+            return null;
+          });
+    }
+    assertFalse(isRouteToLeader);
+  }
+
   private SpannerOptions createSpannerOptions() {
     String endpoint = address.getHostString() + ":" + server.getPort();
     return SpannerOptions.newBuilder()
