set hostId in HostSpec (#782)

aws · Dec 12, 2023 · db4d972 · db4d972
1 parent 9599da2
commit db4d972
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 9 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/HostSpec.java b/wrapper/src/main/java/software/amazon/jdbc/HostSpec.java
@@ -46,15 +46,31 @@ public class HostSpec {
   protected Timestamp lastUpdateTime;
   protected HostAvailabilityStrategy hostAvailabilityStrategy;
 
-  private HostSpec(final String host, final int port, final HostRole role, final HostAvailability availability,
+  private HostSpec(
+      final String host,
+      final int port,
+      final String hostId,
+      final HostRole role,
+      final HostAvailability availability,
       final HostAvailabilityStrategy hostAvailabilityStrategy) {
-    this(host, port, role, availability, DEFAULT_WEIGHT, Timestamp.from(Instant.now()), hostAvailabilityStrategy);
+
+    this(host, port, hostId, role, availability, DEFAULT_WEIGHT,
+        Timestamp.from(Instant.now()), hostAvailabilityStrategy);
   }
 
-  HostSpec(final String host, final int port, final HostRole role, final HostAvailability availability,
-       final long weight, final Timestamp lastUpdateTime, final HostAvailabilityStrategy hostAvailabilityStrategy) {
+  HostSpec(
+      final String host,
+      final int port,
+      final String hostId,
+      final HostRole role,
+      final HostAvailability availability,
+      final long weight,
+      final Timestamp lastUpdateTime,
+      final HostAvailabilityStrategy hostAvailabilityStrategy) {
+
     this.host = host;
     this.port = port;
+    this.hostId = hostId;
     this.availability = availability;
     this.role = role;
     this.allAliases.add(this.asAlias());
@@ -70,7 +86,7 @@ private HostSpec(final String host, final int port, final HostRole role, final H
    * @param role     the role of this host (writer or reader).
    */
   public HostSpec(final HostSpec copyHost, final HostRole role) {
-    this(copyHost.getHost(), copyHost.getPort(), role, copyHost.getAvailability(),
+    this(copyHost.getHost(), copyHost.getPort(), copyHost.getHostId(), role, copyHost.getAvailability(),
         copyHost.getHostAvailabilityStrategy());
   }
 

diff --git a/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java b/wrapper/src/main/java/software/amazon/jdbc/HostSpecBuilder.java
@@ -21,10 +21,10 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.hostavailability.HostAvailabilityStrategy;
-import software.amazon.jdbc.hostavailability.SimpleHostAvailabilityStrategy;
 
 public class HostSpecBuilder {
   private String host;
+  private String hostId;
   private int port = HostSpec.NO_PORT;
   private HostAvailability availability = HostAvailability.AVAILABLE;
   private HostRole role = HostRole.WRITER;
@@ -39,6 +39,7 @@ public HostSpecBuilder(final @NonNull HostAvailabilityStrategy hostAvailabilityS
   public HostSpecBuilder(HostSpecBuilder hostSpecBuilder) {
     this.host = hostSpecBuilder.host;
     this.port = hostSpecBuilder.port;
+    this.hostId = hostSpecBuilder.hostId;
     this.availability = hostSpecBuilder.availability;
     this.role = hostSpecBuilder.role;
     this.weight = hostSpecBuilder.weight;
@@ -56,6 +57,11 @@ public HostSpecBuilder port(int port) {
     return this;
   }
 
+  public HostSpecBuilder hostId(String hostId) {
+    this.hostId = hostId;
+    return this;
+  }
+
   public HostSpecBuilder availability(HostAvailability availability) {
     this.availability = availability;
     return this;
@@ -84,8 +90,8 @@ public HostSpecBuilder lastUpdateTime(Timestamp lastUpdateTime) {
   public HostSpec build() {
     checkHostIsSet();
     setDefaultLastUpdateTime();
-    return new HostSpec(this.host, this.port, this.role, this.availability, this.weight, this.lastUpdateTime,
-        this.hostAvailabilityStrategy);
+    return new HostSpec(this.host, this.port, this.hostId, this.role, this.availability,
+        this.weight, this.lastUpdateTime, this.hostAvailabilityStrategy);
   }
 
   private void checkHostIsSet() {

diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/ConnectionUrlParser.java b/wrapper/src/main/java/software/amazon/jdbc/util/ConnectionUrlParser.java
@@ -27,7 +27,6 @@
 import software.amazon.jdbc.HostRole;
 import software.amazon.jdbc.HostSpec;
 import software.amazon.jdbc.HostSpecBuilder;
-import software.amazon.jdbc.hostavailability.HostAvailabilityStrategyFactory;
 
 public class ConnectionUrlParser {
 
@@ -90,6 +89,8 @@ public static HostSpec parseHostPortPair(final String url, final HostRole role,
 
   private static HostSpec getHostSpec(final String[] hostPortPair, final HostRole hostRole,
       final HostSpecBuilder hostSpecBuilder) {
+    String hostId = rdsUtils.getRdsInstanceId(hostPortPair[0]);
+
     if (hostPortPair.length > 1) {
       final String[] port = hostPortPair[1].split("/");
       int portValue = parsePortAsInt(hostPortPair[1]);
@@ -99,12 +100,14 @@ private static HostSpec getHostSpec(final String[] hostPortPair, final HostRole
       return hostSpecBuilder
           .host(hostPortPair[0])
           .port(portValue)
+          .hostId(hostId)
           .role(hostRole)
           .build();
     }
     return hostSpecBuilder
         .host(hostPortPair[0])
         .port(HostSpec.NO_PORT)
+        .hostId(hostId)
         .role(hostRole)
         .build();
   }

diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java
@@ -18,6 +18,7 @@
 
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 public class RdsUtils {
 
@@ -141,6 +142,8 @@ public class RdsUtils {
       Pattern.compile(
           "^(([0-9A-Fa-f]{1,4}(:[0-9A-Fa-f]{1,4}){0,5})?)"
               + "::(([0-9A-Fa-f]{1,4}(:[0-9A-Fa-f]{1,4}){0,5})?)$");
+
+  private static final String INSTANCE_GROUP = "instance";
   private static final String DNS_GROUP = "dns";
   private static final String DOMAIN_GROUP = "domain";
   private static final String REGION_GROUP = "region";
@@ -178,6 +181,21 @@ public boolean isElbUrl(final String host) {
         && (ELB_PATTERN.matcher(host).find());
   }
 
+  public @Nullable String getRdsInstanceId(final String host) {
+    if (StringUtils.isNullOrEmpty(host)) {
+      return null;
+    }
+    final Matcher matcher = AURORA_INSTANCE_PATTERN.matcher(host);
+    if (matcher.find()) {
+      return matcher.group(INSTANCE_GROUP);
+    }
+    final Matcher matcherChina = AURORA_CHINA_INSTANCE_PATTERN.matcher(host);
+    if (matcherChina.find()) {
+      return matcherChina.group(INSTANCE_GROUP);
+    }
+    return null;
+  }
+
   public String getRdsInstanceHostPattern(final String host) {
     if (StringUtils.isNullOrEmpty(host)) {
       return "?";