fix: connection identification and tracking (aws#943)

Bit-Quill · Apr 1, 2024 · f57a92e · f57a92e
1 parent 1140de8
commit f57a92e
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 17 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/HostSpec.java b/wrapper/src/main/java/software/amazon/jdbc/HostSpec.java
@@ -25,6 +25,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import software.amazon.jdbc.hostavailability.HostAvailability;
 import software.amazon.jdbc.hostavailability.HostAvailabilityStrategy;
+import software.amazon.jdbc.util.StringUtils;
 
 /**
  * An object representing connection info for a given host. Modifiable fields are thread-safe to support sharing this
@@ -150,8 +151,10 @@ public void addAlias(final String... alias) {
     }
 
     Arrays.asList(alias).forEach(x -> {
-      this.aliases.add(x);
-      this.allAliases.add(x);
+      if (!StringUtils.isNullOrEmpty(x)) {
+        this.aliases.add(x);
+        this.allAliases.add(x);
+      }
     });
   }
 
@@ -160,8 +163,10 @@ public void removeAlias(final String... alias) {
       return;
     }
     Arrays.asList(alias).forEach(x -> {
-      this.aliases.remove(x);
-      this.allAliases.remove(x);
+      if (!StringUtils.isNullOrEmpty(x)) {
+        this.aliases.remove(x);
+        this.allAliases.remove(x);
+      }
     });
   }
 

diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/RdsHostListProvider.java
@@ -648,17 +648,38 @@ public HostSpec identifyConnection(Connection connection) throws SQLException {
       if (resultSet.next()) {
         final String instanceName = resultSet.getString(1);
 
-        final List<HostSpec> topology = this.refresh();
+        List<HostSpec> topology = this.refresh(connection);
+
+        boolean isForcedRefresh = false;
+        if (topology == null) {
+          topology = this.forceRefresh(connection);
+          isForcedRefresh = true;
+        }
 
         if (topology == null) {
           return null;
         }
 
-        return topology
+        HostSpec foundHost = topology
             .stream()
             .filter(host -> Objects.equals(instanceName, host.getHostId()))
             .findAny()
             .orElse(null);
+
+        if (foundHost == null && !isForcedRefresh) {
+          topology = this.forceRefresh(connection);
+          if (topology == null) {
+            return null;
+          }
+
+          foundHost = topology
+              .stream()
+              .filter(host -> Objects.equals(instanceName, host.getHostId()))
+              .findAny()
+              .orElse(null);
+        }
+
+        return foundHost;
       }
     } catch (final SQLException e) {
       throw new SQLException(Messages.get("RdsHostListProvider.errorIdentifyConnection"), e);

diff --git a/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java b/wrapper/src/main/java/software/amazon/jdbc/util/RdsUtils.java
@@ -141,7 +141,7 @@ public boolean isRdsDns(final String host) {
   }
 
   public boolean isRdsInstance(final String host) {
-    return getDnsGroup(host) == null;
+    return getDnsGroup(host) == null && isRdsDns(host);
   }
 
   public boolean isRdsProxyDns(final String host) {

diff --git a/wrapper/src/test/java/integration/host/TestEnvironment.java b/wrapper/src/test/java/integration/host/TestEnvironment.java
@@ -494,17 +494,18 @@ private static String getAuroraDbEngineVersion(TestEnvironment env) {
       default:
         throw new NotImplementedException(request.getDatabaseEngine().toString());
     }
-    return findAuroraDbEngineVersion(env, engineName, systemPropertyVersion.toLowerCase());
+    return findAuroraDbEngineVersion(env, engineName, systemPropertyVersion);
   }
 
   private static String findAuroraDbEngineVersion(
       TestEnvironment env,
       String engineName,
       String systemPropertyVersion) {
-    if (systemPropertyVersion == null) {
+
+    if (StringUtils.isNullOrEmpty(systemPropertyVersion)) {
       return env.auroraUtil.getLTSVersion(engineName);
     }
-    switch (systemPropertyVersion) {
+    switch (systemPropertyVersion.toLowerCase()) {
       case "lts":
         return env.auroraUtil.getLTSVersion(engineName);
       case "latest":

diff --git a/wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java b/wrapper/src/test/java/software/amazon/jdbc/hostlistprovider/RdsHostListProviderTest.java
@@ -473,7 +473,8 @@ void testIdentifyConnectionNullTopology() throws SQLException {
 
     when(mockResultSet.next()).thenReturn(true);
     when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
-    when(rdsHostListProvider.refresh(eq(mockConnection))).thenReturn(null);
+    doReturn(null).when(rdsHostListProvider).refresh(mockConnection);
+    doReturn(null).when(rdsHostListProvider).forceRefresh(mockConnection);
 
     assertNull(rdsHostListProvider.identifyConnection(mockConnection));
   }
@@ -492,7 +493,8 @@ void testIdentifyConnectionHostNotInTopology() throws SQLException {
         "jdbc:someprotocol://url"));
     when(mockResultSet.next()).thenReturn(true);
     when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
-    when(rdsHostListProvider.refresh(eq(mockConnection))).thenReturn(cachedTopology);
+    doReturn(cachedTopology).when(rdsHostListProvider).refresh(mockConnection);
+    doReturn(cachedTopology).when(rdsHostListProvider).forceRefresh(mockConnection);
 
     assertNull(rdsHostListProvider.identifyConnection(mockConnection));
   }
@@ -512,7 +514,8 @@ void testIdentifyConnectionHostInTopology() throws SQLException {
         "jdbc:someprotocol://url"));
     when(mockResultSet.next()).thenReturn(true);
     when(mockResultSet.getString(eq(1))).thenReturn("instance-a-1");
-    when(rdsHostListProvider.refresh()).thenReturn(cachedTopology);
+    doReturn(cachedTopology).when(rdsHostListProvider).refresh(mockConnection);
+    doReturn(cachedTopology).when(rdsHostListProvider).forceRefresh(mockConnection);
 
     final HostSpec actual = rdsHostListProvider.identifyConnection(mockConnection);
     assertEquals("instance-a-1.xyz.us-east-2.rds.amazonaws.com", actual.getHost());

diff --git a/.../test/java/software/amazon/jdbc/hostlistprovider/RdsMultiAzDbClusterListProviderTest.java b/.../test/java/software/amazon/jdbc/hostlistprovider/RdsMultiAzDbClusterListProviderTest.java
@@ -46,6 +46,7 @@
 import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
@@ -454,7 +455,8 @@ void testIdentifyConnectionNullTopology() throws SQLException {
 
     when(mockResultSet.next()).thenReturn(true);
     when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
-    when(rdsMazDbClusterHostListProvider.refresh(eq(mockConnection))).thenReturn(null);
+    doReturn(null).when(rdsMazDbClusterHostListProvider).refresh(mockConnection);
+    doReturn(null).when(rdsMazDbClusterHostListProvider).forceRefresh(mockConnection);
 
     assertNull(rdsMazDbClusterHostListProvider.identifyConnection(mockConnection));
   }
@@ -473,7 +475,8 @@ void testIdentifyConnectionHostNotInTopology() throws SQLException {
         "jdbc:someprotocol://url"));
     when(mockResultSet.next()).thenReturn(true);
     when(mockResultSet.getString(eq(1))).thenReturn("instance-1");
-    when(rdsMazDbClusterHostListProvider.refresh(eq(mockConnection))).thenReturn(cachedTopology);
+    doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).refresh(mockConnection);
+    doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).forceRefresh(mockConnection);
 
     assertNull(rdsMazDbClusterHostListProvider.identifyConnection(mockConnection));
   }
@@ -482,18 +485,19 @@ void testIdentifyConnectionHostNotInTopology() throws SQLException {
   void testIdentifyConnectionHostInTopology() throws SQLException {
     final HostSpec expectedHost = new HostSpecBuilder(new SimpleHostAvailabilityStrategy())
         .host("instance-a-1.xyz.us-east-2.rds.amazonaws.com")
+        .hostId("instance-a-1")
         .port(HostSpec.NO_PORT)
         .role(HostRole.WRITER)
         .build();
-    expectedHost.setHostId("instance-a-1");
     final List<HostSpec> cachedTopology = Collections.singletonList(expectedHost);
 
     rdsMazDbClusterHostListProvider = Mockito.spy(getRdsMazDbClusterHostListProvider(
         mockHostListProviderService,
         "jdbc:someprotocol://url"));
     when(mockResultSet.next()).thenReturn(true);
     when(mockResultSet.getString(eq(1))).thenReturn("instance-a-1");
-    when(rdsMazDbClusterHostListProvider.refresh()).thenReturn(cachedTopology);
+    doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).refresh(mockConnection);
+    doReturn(cachedTopology).when(rdsMazDbClusterHostListProvider).forceRefresh(mockConnection);
 
     final HostSpec actual = rdsMazDbClusterHostListProvider.identifyConnection(mockConnection);
     assertEquals("instance-a-1.xyz.us-east-2.rds.amazonaws.com", actual.getHost());