feat: Leader Aware Routing in Connection API (#2308)

* feat: Leader Aware Routing in Connection API

(cherry picked from commit 83ded36)

* Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

Co-authored-by: Knut Olav Løite <[email protected]>

* add verbose documentation

* address review comments

* change assertThat to assertEquals/assertNotEquals within testGetSpanner method

* add credentials url

* add credentials url

---------

Co-authored-by: Knut Olav Løite <[email protected]>

Author: rajatbhatta

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

@@ -172,6 +172,7 @@ public String[] getValidValues() {
   private static final RpcPriority DEFAULT_RPC_PRIORITY = null;
   private static final boolean DEFAULT_RETURN_COMMIT_STATS = false;
   private static final boolean DEFAULT_LENIENT = false;
+  private static final boolean DEFAULT_ROUTE_TO_LEADER = true;
   private static final boolean DEFAULT_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE = false;
   private static final boolean DEFAULT_TRACK_SESSION_LEAKS = true;
   private static final boolean DEFAULT_TRACK_CONNECTION_LEAKS = true;
@@ -186,6 +187,8 @@ public String[] getValidValues() {
   public static final String AUTOCOMMIT_PROPERTY_NAME = "autocommit";
   /** Name of the 'readonly' connection property. */
   public static final String READONLY_PROPERTY_NAME = "readonly";
+  /** Name of the 'routeToLeader' connection property. */
+  public static final String ROUTE_TO_LEADER_PROPERTY_NAME = "routeToLeader";
   /** Name of the 'retry aborts internally' connection property. */
   public static final String RETRY_ABORTS_INTERNALLY_PROPERTY_NAME = "retryAbortsInternally";
   /** Name of the 'credentials' connection property. */
@@ -241,6 +244,10 @@ public String[] getValidValues() {
                       READONLY_PROPERTY_NAME,
                       "Should the connection start in read-only mode (true/false)",
                       DEFAULT_READONLY),
+                  ConnectionProperty.createBooleanProperty(
+                      ROUTE_TO_LEADER_PROPERTY_NAME,
+                      "Should read/write transactions and partitioned DML be routed to leader region (true/false)",
+                      DEFAULT_ROUTE_TO_LEADER),
                   ConnectionProperty.createBooleanProperty(
                       RETRY_ABORTS_INTERNALLY_PROPERTY_NAME,
                       "Should the connection automatically retry Aborted errors (true/false)",
@@ -462,6 +469,8 @@ private boolean isValidUri(String uri) {
      *       created on the emulator if any of them do not yet exist. Any existing instance or
      *       database on the emulator will remain untouched. No other configuration is needed in
      *       order to connect to the emulator than setting this property.
+     *   <li>routeToLeader (boolean): Sets the routeToLeader flag to route requests to leader (true)
+     *       or any region (false) in read/write transactions and Partitioned DML. Default is true.
      * </ul>
      *
      * @param uri The URI of the Spanner database to connect to.
@@ -586,6 +595,7 @@ public static Builder newBuilder() {
 
   private final boolean autocommit;
   private final boolean readOnly;
+  private final boolean routeToLeader;
   private final boolean retryAbortsInternally;
   private final List<StatementExecutionInterceptor> statementExecutionInterceptors;
   private final SpannerOptionsConfigurator configurator;
@@ -678,6 +688,7 @@ private ConnectionOptions(Builder builder) {
 
     this.autocommit = parseAutocommit(this.uri);
     this.readOnly = parseReadOnly(this.uri);
+    this.routeToLeader = parseRouteToLeader(this.uri);
     this.retryAbortsInternally = parseRetryAbortsInternally(this.uri);
     this.statementExecutionInterceptors =
         Collections.unmodifiableList(builder.statementExecutionInterceptors);
@@ -762,6 +773,11 @@ static boolean parseReadOnly(String uri) {
     return value != null ? Boolean.parseBoolean(value) : DEFAULT_READONLY;
   }
 
+  static boolean parseRouteToLeader(String uri) {
+    String value = parseUriProperty(uri, ROUTE_TO_LEADER_PROPERTY_NAME);
+    return value != null ? Boolean.parseBoolean(value) : DEFAULT_ROUTE_TO_LEADER;
+  }
+
   @VisibleForTesting
   static boolean parseRetryAbortsInternally(String uri) {
     String value = parseUriProperty(uri, RETRY_ABORTS_INTERNALLY_PROPERTY_NAME);
@@ -1089,6 +1105,14 @@ public boolean isReadOnly() {
     return readOnly;
   }
 
+  /**
+   * Whether read/write transactions and partitioned DML are preferred to be routed to the leader
+   * region.
+   */
+  public boolean isRouteToLeader() {
+    return routeToLeader;
+  }
+
   /**
    * The initial retryAbortsInternally value for connections created by this {@link
    * ConnectionOptions}

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java

@@ -154,6 +154,7 @@ static class SpannerPoolKey {
     private final boolean usePlainText;
     private final String userAgent;
     private final String databaseRole;
+    private final boolean routeToLeader;
 
     @VisibleForTesting
     static SpannerPoolKey of(ConnectionOptions options) {
@@ -179,6 +180,7 @@ private SpannerPoolKey(ConnectionOptions options) throws IOException {
       this.numChannels = options.getNumChannels();
       this.usePlainText = options.isUsePlainText();
       this.userAgent = options.getUserAgent();
+      this.routeToLeader = options.isRouteToLeader();
     }
 
     @Override
@@ -194,7 +196,8 @@ public boolean equals(Object o) {
           && Objects.equals(this.numChannels, other.numChannels)
           && Objects.equals(this.databaseRole, other.databaseRole)
           && Objects.equals(this.usePlainText, other.usePlainText)
-          && Objects.equals(this.userAgent, other.userAgent);
+          && Objects.equals(this.userAgent, other.userAgent)
+          && Objects.equals(this.routeToLeader, other.routeToLeader);
     }
 
     @Override
@@ -207,7 +210,8 @@ public int hashCode() {
           this.numChannels,
           this.usePlainText,
           this.databaseRole,
-          this.userAgent);
+          this.userAgent,
+          this.routeToLeader);
     }
   }
 
@@ -342,6 +346,9 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
     if (options.getChannelProvider() != null) {
       builder.setChannelProvider(options.getChannelProvider());
     }
+    if (!options.isRouteToLeader()) {
+      builder.disableLeaderAwareRouting();
+    }
     if (key.usePlainText) {
       // Credentials may not be sent over a plain text channel.
       builder.setCredentials(NoCredentials.getInstance());

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java

@@ -49,6 +49,9 @@ public class ConnectionOptionsTest {
   private static final String FILE_TEST_PATH =
       ConnectionOptionsTest.class.getResource("test-key.json").getFile();
   private static final String DEFAULT_HOST = "https://spanner.googleapis.com";
+  private static final String TEST_PROJECT = "test-project-123";
+  private static final String TEST_INSTANCE = "test-instance-123";
+  private static final String TEST_DATABASE = "test-database-123";
 
   @Test
   public void testBuildWithURIWithDots() {
@@ -149,6 +152,27 @@ public void testBuildWithAutoConfigEmulator() {
     assertTrue(options.isUsePlainText());
   }
 
+  @Test
+  public void testBuildWithRouteToLeader() {
+    final String BASE_URI =
+        "cloudspanner:/projects/test-project-123/instances/test-instance-123/databases/test-database-123";
+    ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
+    builder.setUri(BASE_URI + "?routeToLeader=false");
+    builder.setCredentialsUrl(FILE_TEST_PATH);
+    ConnectionOptions options = builder.build();
+    assertEquals(options.getHost(), DEFAULT_HOST);
+    assertEquals(options.getProjectId(), TEST_PROJECT);
+    assertEquals(options.getInstanceId(), TEST_INSTANCE);
+    assertEquals(options.getDatabaseName(), TEST_DATABASE);
+    assertFalse(options.isRouteToLeader());
+
+    // Test for default behavior for routeToLeader property.
+    builder = ConnectionOptions.newBuilder().setUri(BASE_URI);
+    builder.setCredentialsUrl(FILE_TEST_PATH);
+    options = builder.build();
+    assertTrue(options.isRouteToLeader());
+  }
+
   @Test
   public void testBuildWithAutoConfigEmulatorAndHost() {
     ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java

@@ -65,6 +65,8 @@ public class SpannerPoolTest {
 
   private ConnectionOptions options5 = mock(ConnectionOptions.class);
   private ConnectionOptions options6 = mock(ConnectionOptions.class);
+  private ConnectionOptions options7 = mock(ConnectionOptions.class);
+  private ConnectionOptions options8 = mock(ConnectionOptions.class);
 
   private SpannerPool createSubjectAndMocks() {
     return createSubjectAndMocks(0L, Ticker.systemTicker());
@@ -93,6 +95,10 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
     // ConnectionOptions with no specific credentials.
     when(options5.getProjectId()).thenReturn("test-project-3");
     when(options6.getProjectId()).thenReturn("test-project-3");
+    when(options7.getProjectId()).thenReturn("test-project-3");
+    when(options7.isRouteToLeader()).thenReturn(true);
+    when(options8.getProjectId()).thenReturn("test-project-3");
+    when(options8.isRouteToLeader()).thenReturn(false);
 
     return pool;
   }
@@ -111,40 +117,43 @@ public void testGetSpanner() {
     // assert equal
     spanner1 = pool.getSpanner(options1, connection1);
     spanner2 = pool.getSpanner(options1, connection2);
-    assertThat(spanner1).isEqualTo(spanner2);
+    assertEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options2, connection1);
     spanner2 = pool.getSpanner(options2, connection2);
-    assertThat(spanner1).isEqualTo(spanner2);
+    assertEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options3, connection1);
     spanner2 = pool.getSpanner(options3, connection2);
-    assertThat(spanner1).isEqualTo(spanner2);
+    assertEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options4, connection1);
     spanner2 = pool.getSpanner(options4, connection2);
-    assertThat(spanner1).isEqualTo(spanner2);
+    assertEquals(spanner1, spanner2);
     // Options 5 and 6 both use default credentials.
     spanner1 = pool.getSpanner(options5, connection1);
     spanner2 = pool.getSpanner(options6, connection2);
-    assertThat(spanner1).isEqualTo(spanner2);
+    assertEquals(spanner1, spanner2);
 
     // assert not equal
     spanner1 = pool.getSpanner(options1, connection1);
     spanner2 = pool.getSpanner(options2, connection2);
-    assertThat(spanner1).isNotEqualTo(spanner2);
+    assertNotEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options1, connection1);
     spanner2 = pool.getSpanner(options3, connection2);
-    assertThat(spanner1).isNotEqualTo(spanner2);
+    assertNotEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options1, connection1);
     spanner2 = pool.getSpanner(options4, connection2);
-    assertThat(spanner1).isNotEqualTo(spanner2);
+    assertNotEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options2, connection1);
     spanner2 = pool.getSpanner(options3, connection2);
-    assertThat(spanner1).isNotEqualTo(spanner2);
+    assertNotEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options2, connection1);
     spanner2 = pool.getSpanner(options4, connection2);
-    assertThat(spanner1).isNotEqualTo(spanner2);
+    assertNotEquals(spanner1, spanner2);
     spanner1 = pool.getSpanner(options3, connection1);
     spanner2 = pool.getSpanner(options4, connection2);
-    assertThat(spanner1).isNotEqualTo(spanner2);
+    assertNotEquals(spanner1, spanner2);
+    spanner1 = pool.getSpanner(options7, connection1);
+    spanner2 = pool.getSpanner(options8, connection2);
+    assertNotEquals(spanner1, spanner2);
   }
 
   @Test
@@ -460,14 +469,30 @@ public void testSpannerPoolKeyEquality() {
             .setUri("cloudspanner:/projects/p/instances/i/databases/d")
             .setCredentials(NoCredentials.getInstance())
             .build();
+    // Not passing in routeToLeader in Connection URI is equivalent to passing it as true,
+    // as routeToLeader is true by default.
+    ConnectionOptions options4 =
+        ConnectionOptions.newBuilder()
+            .setUri("cloudspanner:/projects/p/instances/i/databases/d?routeToLeader=true")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    ConnectionOptions options5 =
+        ConnectionOptions.newBuilder()
+            .setUri("cloudspanner:/projects/p/instances/i/databases/d?routeToLeader=false")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
 
     SpannerPoolKey key1 = SpannerPoolKey.of(options1);
     SpannerPoolKey key2 = SpannerPoolKey.of(options2);
     SpannerPoolKey key3 = SpannerPoolKey.of(options3);
+    SpannerPoolKey key4 = SpannerPoolKey.of(options4);
+    SpannerPoolKey key5 = SpannerPoolKey.of(options5);
 
     assertNotEquals(key1, key2);
     assertEquals(key2, key3);
     assertNotEquals(key1, key3);
     assertNotEquals(key1, new Object());
+    assertEquals(key3, key4);
+    assertNotEquals(key4, key5);
   }
 }