refactor: host alias logic in wrapper

wrapper/src/main/java/software/amazon/jdbc/Driver.java

@@ -91,6 +91,8 @@ public Connection connect(final String url, final Properties info) throws SQLExc
       return null;
     }
 
+    LOGGER.finest("Opening connection to " + url);
+
     final String driverUrl = url.replaceFirst(PROTOCOL_PREFIX, "jdbc:");
     final java.sql.Driver driver = DriverManager.getDriver(driverUrl);
 

wrapper/src/main/java/software/amazon/jdbc/HostListProvider.java

@@ -39,4 +39,6 @@ public interface HostListProvider {
    *                      determine the host role
    */
   HostRole getHostRole(Connection connection) throws SQLException;
+
+  HostSpec identifyConnection(Connection connection) throws SQLException;
 }

wrapper/src/main/java/software/amazon/jdbc/HostSpec.java

@@ -38,6 +38,7 @@ public class HostSpec {
   protected Set<String> aliases = ConcurrentHashMap.newKeySet();
   protected Set<String> allAliases = ConcurrentHashMap.newKeySet();
   protected long weight; // Greater or equal 0. Lesser the weight, the healthier node.
+  protected String hostId;
 
   public HostSpec(final String host) {
     this.host = host;
@@ -148,6 +149,12 @@ public void removeAlias(final String... alias) {
     });
   }
 
+  public void resetAliases() {
+    this.aliases.clear();
+    this.allAliases.clear();
+    this.allAliases.add(this.asAlias());
+  }
+
   public String getUrl() {
     String url = isPortSpecified() ? host + ":" + port : host;
     if (!url.endsWith("/")) {
@@ -156,6 +163,14 @@ public String getUrl() {
     return url;
   }
 
+  public String getHostId() {
+    return hostId;
+  }
+
+  public void setHostId(String hostId) {
+    this.hostId = hostId;
+  }
+
   public String asAlias() {
     return isPortSpecified() ? host + ":" + port : host;
   }

wrapper/src/main/java/software/amazon/jdbc/PluginService.java

@@ -151,4 +151,8 @@ HostSpec getHostSpecByStrategy(HostRole role, String strategy)
   Dialect getDialect();
 
   void updateDialect(final @NonNull Connection connection) throws SQLException;
+
+  HostSpec identifyConnection(final Connection connection) throws SQLException;
+
+  void fillAliases(final Connection connection, final HostSpec hostSpec) throws SQLException;
 }

wrapper/src/main/java/software/amazon/jdbc/PluginServiceImpl.java

@@ -17,13 +17,16 @@
 package software.amazon.jdbc;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -35,6 +38,7 @@
 import software.amazon.jdbc.dialect.Dialect;
 import software.amazon.jdbc.dialect.DialectManager;
 import software.amazon.jdbc.dialect.DialectProvider;
+import software.amazon.jdbc.dialect.TopologyAwareDatabaseCluster;
 import software.amazon.jdbc.exceptions.ExceptionManager;
 import software.amazon.jdbc.hostlistprovider.StaticHostListProvider;
 import software.amazon.jdbc.util.CacheMap;
@@ -334,7 +338,7 @@ public HostListProvider getHostListProvider() {
   @Override
   public void refreshHostList() throws SQLException {
     final List<HostSpec> updatedHostList = this.getHostListProvider().refresh();
-    if (updatedHostList != null) {
+    if (!Objects.equals(updatedHostList, this.hosts)) {
       updateHostAvailability(updatedHostList);
       setNodeList(this.hosts, updatedHostList);
     }
@@ -343,7 +347,7 @@ public void refreshHostList() throws SQLException {
   @Override
   public void refreshHostList(final Connection connection) throws SQLException {
     final List<HostSpec> updatedHostList = this.getHostListProvider().refresh(connection);
-    if (updatedHostList != null) {
+    if (!Objects.equals(updatedHostList, this.hosts)) {
       updateHostAvailability(updatedHostList);
       setNodeList(this.hosts, updatedHostList);
     }
@@ -489,4 +493,44 @@ public void updateDialect(final @NonNull Connection connection) throws SQLExcept
         connection);
   }
 
+  @Override
+  public HostSpec identifyConnection(Connection connection) throws SQLException {
+    if (!(this.getDialect() instanceof TopologyAwareDatabaseCluster)) {
+      return null;
+    }
+
+    return this.hostListProvider.identifyConnection(connection);
+  }
+
+  @Override
+  public void fillAliases(Connection connection, HostSpec hostSpec) throws SQLException {
+    if (hostSpec == null) {
+      return;
+    }
+
+    if (!hostSpec.getAliases().isEmpty()) {
+      LOGGER.finest(() -> Messages.get("PluginServiceImpl.nonEmptyAliases", new Object[] {hostSpec.getAliases()}));
+      return;
+    }
+
+    hostSpec.addAlias(hostSpec.asAlias());
+
+    // Add the host name and port, this host name is usually the internal IP address.
+    try (final Statement stmt = connection.createStatement()) {
+      try (final ResultSet rs = stmt.executeQuery(this.getDialect().getHostAliasQuery())) {
+        while (rs.next()) {
+          hostSpec.addAlias(rs.getString(1));
+        }
+      }
+    } catch (final SQLException sqlException) {
+      // log and ignore
+      LOGGER.finest(() -> Messages.get("PluginServiceImpl.failedToRetrieveHostPort"));
+    }
+
+    // Add the instance endpoint if the current connection is associated with a topology aware database cluster.
+    final HostSpec host = this.identifyConnection(connection);
+    if (host != null) {
+      hostSpec.addAlias(host.asAliases().toArray(new String[] {}));
+    }
+  }
 }

wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/AuroraHostListProvider.java

@@ -27,6 +27,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map.Entry;
+import java.util.Objects;
 import java.util.Properties;
 import java.util.Set;
 import java.util.UUID;
@@ -85,7 +86,6 @@ public class AuroraHostListProvider implements DynamicHostListProvider {
       : TimeUnit.MILLISECONDS.toNanos(30000);
   private final long suggestedClusterIdRefreshRateNano = TimeUnit.MINUTES.toNanos(10);
   private List<HostSpec> hostList = new ArrayList<>();
-  private List<HostSpec> lastReturnedHostList;
   private List<HostSpec> initialHostList = new ArrayList<>();
   private HostSpec initialHostSpec;
 
@@ -273,7 +273,7 @@ private ClusterSuggestedResult getSuggestedClusterId(final String url) {
       for (final HostSpec host : hosts) {
         if (host.getUrl().equals(url)) {
           LOGGER.finest(() -> Messages.get("AuroraHostListProvider.suggestedClusterId",
-              new Object[]{key, url}));
+              new Object[] {key, url}));
           return new ClusterSuggestedResult(key, isPrimaryCluster);
         }
       }
@@ -398,8 +398,12 @@ private HostSpec createHost(final ResultSet resultSet) throws SQLException {
     // Calculate weight based on node lag in time and CPU utilization.
     final long weight = Math.round(nodeLag) * 100L + Math.round(cpuUtilization);
 
-    hostName = hostName == null ? "?" : hostName;
-    final String endpoint = getHostEndpoint(hostName);
+    return createHost(hostName, isWriter, weight);
+  }
+
+  private HostSpec createHost(String host, final boolean isWriter, final long weight) {
+    host = host == null ? "?" : host;
+    final String endpoint = getHostEndpoint(host);
     final int port = this.clusterInstanceTemplate.isPortSpecified()
         ? this.clusterInstanceTemplate.getPort()
         : this.initialHostSpec.getPort();
@@ -410,7 +414,8 @@ private HostSpec createHost(final ResultSet resultSet) throws SQLException {
         isWriter ? HostRole.WRITER : HostRole.READER,
         HostAvailability.AVAILABLE,
         weight);
-    hostSpec.addAlias(hostName);
+    hostSpec.addAlias(host);
+    hostSpec.setHostId(host);
     return hostSpec;
   }
 
@@ -466,12 +471,7 @@ public List<HostSpec> refresh(final Connection connection) throws SQLException {
     final FetchTopologyResult results = getTopology(currentConnection, false);
     LOGGER.finest(() -> Utils.logTopology(results.hosts));
 
-    if (results.isCachedData && this.lastReturnedHostList == results.hosts) {
-      return null; // no topology update
-    }
-
     this.hostList = results.hosts;
-    this.lastReturnedHostList = this.hostList;
     return Collections.unmodifiableList(hostList);
   }
 
@@ -490,7 +490,6 @@ public List<HostSpec> forceRefresh(final Connection connection) throws SQLExcept
     final FetchTopologyResult results = getTopology(currentConnection, true);
     LOGGER.finest(() -> Utils.logTopology(results.hosts));
     this.hostList = results.hosts;
-    this.lastReturnedHostList = this.hostList;
     return Collections.unmodifiableList(this.hostList);
   }
 
@@ -577,6 +576,7 @@ public FetchTopologyResult(final boolean isCachedData, final List<HostSpec> host
   }
 
   static class ClusterSuggestedResult {
+
     public String clusterId;
     public boolean isPrimaryClusterId;
 
@@ -588,18 +588,10 @@ public ClusterSuggestedResult(final String clusterId, final boolean isPrimaryClu
 
   @Override
   public HostRole getHostRole(Connection conn) throws SQLException {
-    if (this.topologyAwareDialect == null) {
-      Dialect dialect = this.hostListProviderService.getDialect();
-      if (!(dialect instanceof TopologyAwareDatabaseCluster)) {
-        throw new SQLException(
-            Messages.get("AuroraHostListProvider.invalidDialectForGetHostRole",
-                new Object[]{dialect}));
-      }
-      this.topologyAwareDialect = (TopologyAwareDatabaseCluster) this.hostListProviderService.getDialect();
-    }
-
     try (final Statement stmt = conn.createStatement();
-         final ResultSet rs = stmt.executeQuery(this.topologyAwareDialect.getIsReaderQuery())) {
+        final ResultSet rs = stmt.executeQuery(
+            getTopologyAwareDialect("AuroraHostListProvider.invalidDialectForGetHostRole")
+                .getIsReaderQuery())) {
       if (rs.next()) {
         boolean isReader = rs.getBoolean(1);
         return isReader ? HostRole.READER : HostRole.WRITER;
@@ -610,4 +602,45 @@ public HostRole getHostRole(Connection conn) throws SQLException {
 
     throw new SQLException(Messages.get("AuroraHostListProvider.errorGettingHostRole"));
   }
+
+  @Override
+  public HostSpec identifyConnection(Connection connection) throws SQLException {
+    try (final Statement stmt = connection.createStatement();
+        final ResultSet resultSet = stmt.executeQuery(
+            getTopologyAwareDialect("AuroraHostListProvider.invalidDialectForIdentifyConnection")
+                .getNodeIdQuery())) {
+      if (resultSet.next()) {
+        final String instanceName = resultSet.getString(1);
+
+        final List<HostSpec> topology = this.refresh();
+
+        if (topology == null) {
+          return null;
+        }
+
+        return topology
+            .stream()
+            .filter(host -> Objects.equals(instanceName, host.getHostId()))
+            .findAny()
+            .orElse(null);
+      }
+    } catch (final SQLException e) {
+      throw new SQLException(Messages.get("AuroraHostListProvider.errorIdentifyConnection"), e);
+    }
+
+    throw new SQLException(Messages.get("AuroraHostListProvider.errorIdentifyConnection"));
+  }
+
+  private TopologyAwareDatabaseCluster getTopologyAwareDialect(String exceptionMessageIdentifier) throws SQLException {
+    if (this.topologyAwareDialect == null) {
+      Dialect dialect = this.hostListProviderService.getDialect();
+      if (!(dialect instanceof TopologyAwareDatabaseCluster)) {
+        throw new SQLException(
+            Messages.get(exceptionMessageIdentifier,
+                new Object[] {dialect}));
+      }
+      this.topologyAwareDialect = (TopologyAwareDatabaseCluster) this.hostListProviderService.getDialect();
+    }
+    return this.topologyAwareDialect;
+  }
 }

wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/ConnectionStringHostListProvider.java

@@ -17,10 +17,13 @@
 package software.amazon.jdbc.hostlistprovider;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Properties;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.AwsWrapperProperty;
@@ -109,4 +112,30 @@ public List<HostSpec> forceRefresh(final Connection connection) throws SQLExcept
   public HostRole getHostRole(Connection connection) {
     throw new UnsupportedOperationException("ConnectionStringHostListProvider does not support getHostRole");
   }
+
+  @Override
+  public HostSpec identifyConnection(Connection connection) throws SQLException {
+    try (final Statement stmt = connection.createStatement();
+        final ResultSet resultSet = stmt.executeQuery(this.hostListProviderService.getDialect().getHostAliasQuery())) {
+      if (resultSet.next()) {
+        final String instance = resultSet.getString(1);
+
+        final List<HostSpec> topology = this.refresh(connection);
+
+        if (topology == null) {
+          return null;
+        }
+
+        return topology
+            .stream()
+            .filter(host -> Objects.equals(instance, host.getHostId()))
+            .findAny()
+            .orElse(null);
+      }
+    } catch (final SQLException e) {
+      throw new SQLException(Messages.get("ConnectionStringHostListProvider.errorIdentifyConnection"), e);
+    }
+
+    throw new SQLException(Messages.get("ConnectionStringHostListProvider.errorIdentifyConnection"));
+  }
 }