From 4e60c463a6c581fba8a54e1caf5a65015f83707f Mon Sep 17 00:00:00 2001
From: sergiyvamz <svvoshch@amazon.com>
Date: Mon, 10 Mar 2025 12:19:54 -0700
Subject: [PATCH] fix issue with non-cluster database dialects and/or custom
 domains

---
 .../main/java/software/amazon/jdbc/PluginService.java  |  2 +-
 .../ConnectionStringHostListProvider.java              | 10 +++++-----
 .../jdbc/plugin/AuroraConnectionTrackerPlugin.java     |  2 +-
 .../plugin/AuroraInitialConnectionStrategyPlugin.java  |  8 +++++++-
 4 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java
index 531abdfb5..c23be9690 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/PluginService.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/PluginService.java
@@ -227,7 +227,7 @@ Connection forceConnect(
 
   void updateDialect(final @NonNull Connection connection) throws SQLException;
 
-  HostSpec identifyConnection(final Connection connection) throws SQLException;
+  @Nullable HostSpec identifyConnection(final Connection connection) throws SQLException;
 
   void fillAliases(final Connection connection, final HostSpec hostSpec) throws SQLException;
 
diff --git a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/ConnectionStringHostListProvider.java b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/ConnectionStringHostListProvider.java
index 328522f27..ecab0fc58 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/ConnectionStringHostListProvider.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/hostlistprovider/ConnectionStringHostListProvider.java
@@ -17,14 +17,12 @@
 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 java.util.logging.Logger;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.jdbc.AwsWrapperProperty;
 import software.amazon.jdbc.HostListProviderService;
@@ -35,6 +33,8 @@
 
 public class ConnectionStringHostListProvider implements StaticHostListProvider {
 
+  private static final Logger LOGGER = Logger.getLogger(ConnectionStringHostListProvider.class.getName());
+
   final List<HostSpec> hostList = new ArrayList<>();
   Properties properties;
   private boolean isInitialized = false;
@@ -115,8 +115,8 @@ public HostRole getHostRole(Connection connection) {
 
   @Override
   public HostSpec identifyConnection(Connection connection) throws SQLException {
-    throw new UnsupportedOperationException(
-        Messages.get("ConnectionStringHostListProvider.unsupportedIdentifyConnection"));
+    LOGGER.finest(Messages.get("ConnectionStringHostListProvider.unsupportedIdentifyConnection"));
+    return null;
   }
 
   @Override
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java
index 863df4ee5..97856e39b 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraConnectionTrackerPlugin.java
@@ -96,7 +96,7 @@ public Connection connect(final String driverProtocol, final HostSpec hostSpec,
 
     if (conn != null) {
       final RdsUrlType type = this.rdsHelper.identifyRdsType(hostSpec.getHost());
-      if (type.isRdsCluster() || type == RdsUrlType.OTHER) {
+      if (type.isRdsCluster() || type == RdsUrlType.OTHER || type == RdsUrlType.IP_ADDRESS) {
         hostSpec.resetAliases();
         this.pluginService.fillAliases(conn, hostSpec);
       }
diff --git a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java
index c2815d1b1..7b1eedbad 100644
--- a/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java
+++ b/wrapper/src/main/java/software/amazon/jdbc/plugin/AuroraInitialConnectionStrategyPlugin.java
@@ -196,7 +196,7 @@ private Connection getVerifiedWriterConnection(
           this.pluginService.forceRefreshHostList(writerCandidateConn);
           writerCandidate = this.pluginService.identifyConnection(writerCandidateConn);
 
-          if (writerCandidate.getRole() != HostRole.WRITER) {
+          if (writerCandidate == null || writerCandidate.getRole() != HostRole.WRITER) {
             // Shouldn't be here. But let's try again.
             this.closeConnection(writerCandidateConn);
             this.delay(retryDelayMs);
@@ -273,6 +273,12 @@ private Connection getVerifiedReaderConnection(
           this.pluginService.forceRefreshHostList(readerCandidateConn);
           readerCandidate = this.pluginService.identifyConnection(readerCandidateConn);
 
+          if (readerCandidate == null) {
+            this.closeConnection(readerCandidateConn);
+            this.delay(retryDelayMs);
+            continue;
+          }
+
           if (readerCandidate.getRole() != HostRole.READER) {
             if (this.hasNoReaders()) {
               // It seems that cluster has no readers. Simulate Aurora reader cluster endpoint logic