Improve Multi-AZ cluster detection (aws#824)

Bit-Quill · Jan 10, 2024 · 9e074d1 · 9e074d1
1 parent 3ff1be1
commit 9e074d1
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 5 deletions.
diff --git a/wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java b/wrapper/src/main/java/software/amazon/jdbc/dialect/DialectManager.java
@@ -66,7 +66,12 @@ public class DialectManager implements DialectProvider {
         }
       };
 
-  protected static final long ENDPOINT_CACHE_EXPIRATION = TimeUnit.MINUTES.toNanos(30);
+  /**
+   * In order to simplify dialect detection, there's an internal host-to-dialect cache.
+   * The cache contains host endpoints and identified dialect. Cache expiration time
+   * is defined by the variable below.
+   */
+  protected static final long ENDPOINT_CACHE_EXPIRATION = TimeUnit.HOURS.toNanos(24);
 
   // Map of host name, or url, by dialect code.
   protected static final CacheMap<String, String> knownEndpointDialects = new CacheMap<>();

diff --git a/wrapper/src/main/java/software/amazon/jdbc/dialect/RdsMultiAzDbClusterMysqlDialect.java b/wrapper/src/main/java/software/amazon/jdbc/dialect/RdsMultiAzDbClusterMysqlDialect.java
@@ -31,6 +31,10 @@ public class RdsMultiAzDbClusterMysqlDialect extends MysqlDialect {
 
   private static final String TOPOLOGY_QUERY = "SELECT id, endpoint, port FROM mysql.rds_topology";
 
+  private static final String TOPOLOGY_TABLE_EXIST_QUERY =
+      "SELECT 1 AS tmp FROM information_schema.tables WHERE"
+      + " table_schema = 'mysql' AND table_name = 'rds_topology'";
+
   private static final String FETCH_WRITER_NODE_QUERY = "SHOW REPLICA STATUS";
 
   private static final String FETCH_WRITER_NODE_QUERY_COLUMN_NAME = "Source_Server_Id";
@@ -44,8 +48,18 @@ public boolean isDialect(final Connection connection) {
     ResultSet rs = null;
     try {
       stmt = connection.createStatement();
-      rs = stmt.executeQuery(TOPOLOGY_QUERY);
-      return rs.next();
+      rs = stmt.executeQuery(TOPOLOGY_TABLE_EXIST_QUERY);
+
+      if (rs.next()) {
+        rs.close();
+        stmt.close();
+
+        stmt = connection.createStatement();
+        rs = stmt.executeQuery(TOPOLOGY_QUERY);
+
+        return rs.next();
+      }
+      return false;
     } catch (final SQLException ex) {
       // ignore
     } finally {

diff --git a/wrapper/src/main/java/software/amazon/jdbc/dialect/RdsMultiAzDbClusterPgDialect.java b/wrapper/src/main/java/software/amazon/jdbc/dialect/RdsMultiAzDbClusterPgDialect.java
@@ -33,6 +33,10 @@ public class RdsMultiAzDbClusterPgDialect extends PgDialect {
   private static final String TOPOLOGY_QUERY =
       "SELECT id, endpoint, port FROM rds_tools.show_topology('aws_jdbc_driver-" + DriverInfo.DRIVER_VERSION + "')";
 
+  private static final String WRITER_NODE_FUNC_EXIST_QUERY =
+      "SELECT 1 AS tmp FROM information_schema.routines"
+          + " WHERE routine_schema='rds_tools' AND routine_name='multi_az_db_cluster_source_dbi_resource_id'";
+
   private static final String FETCH_WRITER_NODE_QUERY =
       "SELECT multi_az_db_cluster_source_dbi_resource_id FROM rds_tools.multi_az_db_cluster_source_dbi_resource_id()";
 
@@ -55,8 +59,18 @@ public boolean isDialect(final Connection connection) {
     ResultSet rs = null;
     try {
       stmt = connection.createStatement();
-      rs = stmt.executeQuery(FETCH_WRITER_NODE_QUERY);
-      return rs.next();
+      rs = stmt.executeQuery(WRITER_NODE_FUNC_EXIST_QUERY);
+
+      if (rs.next()) {
+        rs.close();
+        stmt.close();
+
+        stmt = connection.createStatement();
+        rs = stmt.executeQuery(FETCH_WRITER_NODE_QUERY);
+
+        return rs.next();
+      }
+      return false;
     } catch (final SQLException ex) {
       // ignore
     } finally {