fix: connection tracker update after failover to new connection

common/lib/host_list_provider/connection_string_host_list_provider.ts

@@ -81,7 +81,7 @@ export class ConnectionStringHostListProvider implements StaticHostListProvider
     return Promise.resolve(this.hostList);
   }
 
-  getHostRole(client: AwsClient): Promise<HostRole> {
+  getHostRole(client: ClientWrapper): Promise<HostRole> {
     throw new AwsWrapperError("ConnectionStringHostListProvider does not support getHostRole.");
   }
 

common/lib/host_list_provider/host_list_provider.ts

@@ -33,7 +33,7 @@ export interface HostListProvider {
 
   forceRefresh(client: ClientWrapper): Promise<HostInfo[]>;
 
-  getHostRole(client: AwsClient, dialect: DatabaseDialect): Promise<HostRole>;
+  getHostRole(client: ClientWrapper, dialect: DatabaseDialect): Promise<HostRole>;
 
   identifyConnection(targetClient: ClientWrapper, dialect: DatabaseDialect): Promise<HostInfo | null>;
 

common/lib/host_list_provider/rds_host_list_provider.ts

@@ -137,13 +137,13 @@ export class RdsHostListProvider implements DynamicHostListProvider {
     throw new AwsWrapperError("Could not retrieve targetClient.");
   }
 
-  async getHostRole(client: AwsClient, dialect: DatabaseDialect): Promise<HostRole> {
+  async getHostRole(client: ClientWrapper, dialect: DatabaseDialect): Promise<HostRole> {
     if (!this.isTopologyAwareDatabaseDialect(dialect)) {
       throw new TypeError(Messages.get("RdsHostListProvider.incorrectDialect"));
     }
 
-    if (client.targetClient) {
-      return dialect.getHostRole(client.targetClient);
+    if (client) {
+      return await dialect.getHostRole(client);
     } else {
       throw new AwsWrapperError(Messages.get("AwsClient targetClient not defined."));
     }

common/lib/pg_client_wrapper.ts

@@ -17,6 +17,7 @@
 import { ClientWrapper } from "./client_wrapper";
 import { HostInfo } from "./host_info";
 import { uniqueId } from "../logutils";
+import { logTopology } from "./utils/utils";
 
 /*
 This an internal wrapper class for a target community driver client created by the NodePostgresPgDriverDialect.

common/lib/plugin_service.ts

@@ -40,7 +40,7 @@ import { ClientWrapper } from "./client_wrapper";
 import { logger } from "../logutils";
 import { Messages } from "./utils/messages";
 import { DatabaseDialectCodes } from "./database_dialect/database_dialect_codes";
-import { getWriter } from "./utils/utils";
+import { getWriter, logTopology } from "./utils/utils";
 import { TelemetryFactory } from "./utils/telemetry/telemetry_factory";
 import { DriverDialect } from "./driver_dialect/driver_dialect";
 import { ConfigurationProfile } from "./profile/configuration_profile";

common/lib/plugins/connection_tracker/aurora_connection_tracker_plugin.ts

@@ -26,6 +26,8 @@ import { FailoverError } from "../../utils/errors";
 import { HostChangeOptions } from "../../host_change_options";
 import { HostRole } from "../../host_role";
 import { OpenedConnectionTracker } from "./opened_connection_tracker";
+import { logger } from "../../../logutils";
+import { logTopology } from "../../utils/utils";
 
 export class AuroraConnectionTrackerPlugin extends AbstractConnectionPlugin implements CanReleaseResources {
   private static readonly subscribedMethods = new Set<string>(["notifyHostListChanged"].concat(SubscribedMethodHelper.NETWORK_BOUND_METHODS));
@@ -100,7 +102,6 @@ export class AuroraConnectionTrackerPlugin extends AbstractConnectionPlugin impl
 
   private async checkWriterChanged(): Promise<void> {
     const hostInfoAfterFailover = this.getWriter(this.pluginService.getHosts());
-
     if (this.currentWriter === null) {
       this.currentWriter = hostInfoAfterFailover;
       this.needUpdateCurrentWriter = false;
@@ -127,10 +128,12 @@ export class AuroraConnectionTrackerPlugin extends AbstractConnectionPlugin impl
   async notifyHostListChanged(changes: Map<string, Set<HostChangeOptions>>): Promise<void> {
     for (const [key, _] of changes.entries()) {
       const hostChanges = changes.get(key);
+
       if (hostChanges) {
         if (hostChanges.has(HostChangeOptions.PROMOTED_TO_READER)) {
           await this.tracker.invalidateAllConnectionsMultipleHosts(key);
         }
+
         if (hostChanges.has(HostChangeOptions.PROMOTED_TO_WRITER)) {
           this.needUpdateCurrentWriter = true;
         }

common/lib/plugins/connection_tracker/opened_connection_tracker.ts

@@ -21,6 +21,7 @@ import { logger } from "../../../logutils";
 import { MapUtils } from "../../utils/map_utils";
 import { Messages } from "../../utils/messages";
 import { PluginService } from "../../plugin_service";
+import { logTopology } from "../../utils/utils";
 
 export class OpenedConnectionTracker {
   static readonly openedConnections: Map<string, Array<WeakRef<ClientWrapper>>> = new Map<string, Array<WeakRef<ClientWrapper>>>();

common/lib/plugins/failover/failover_plugin.ts

@@ -48,6 +48,7 @@ export class FailoverPlugin extends AbstractConnectionPlugin {
   private static readonly TELEMETRY_WRITER_FAILOVER = "failover to writer instance";
   private static readonly TELEMETRY_READER_FAILOVER = "failover to replica";
   private static readonly METHOD_END = "end";
+
   private static readonly subscribedMethods: Set<string> = new Set([
     "initHostProvider",
     "connect",

common/lib/plugins/failover/reader_failover_handler.ts

@@ -322,6 +322,21 @@ class ConnectionAttemptTask {
     );
     try {
       this.targetClient = await this.pluginService.forceConnect(this.newHost, copy);
+
+      // ensure that new connection is a connection to a reader node
+      try {
+        // eslint-disable-next-line no-constant-condition
+        while (1) {
+          await this.pluginService.forceRefreshHostList();
+          if ((await this.pluginService.getHostRole(this.targetClient)) === HostRole.READER) {
+            break;
+          }
+          await sleep(1000);
+        }
+      } catch (error: any) {
+        logger.debug(Messages.get("ClusterAwareReaderFailoverHandler.errorGettingHostRole", error.message));
+      }
+
       this.pluginService.setAvailability(this.newHost.allAliases, HostAvailability.AVAILABLE);
       logger.info(Messages.get("ClusterAwareReaderFailoverHandler.successfulReaderConnection", this.newHost.host));
       if (this.taskHandler.getSelectedConnectionAttemptTask(this.failoverTaskId) === -1) {

common/lib/utils/locales/en.json

@@ -48,6 +48,7 @@
   "ClusterAwareWriterFailoverHandler.standaloneHost": "[TaskB] Host %s is not yet connected to a cluster. The cluster is still being reconfigured.",
   "ClusterAwareWriterFailoverHandler.taskBAttemptConnectionToNewWriter": "[TaskB] Trying to connect to a new writer: '%s'",
   "ClusterAwareWriterFailoverHandler.alreadyWriter": "Current reader connection is actually a new writer connection.",
+  "ClusterAwareReaderFailoverHandler.errorGettingHostRole": "An error occurred while trying to determine the role of the reader candidate: ",
   "Failover.TransactionResolutionUnknownError": "Transaction resolution unknown. Please re-configure session state if required and try restarting the transaction.",
   "Failover.connectionChangedError": "The active SQL connection has changed due to a connection failure. Please re-configure session state if required.",
   "Failover.parameterValue": "%s = %s",
@@ -193,5 +194,6 @@
   "ConfigurationProfileBuilder.notFound": "Configuration profile '%s' not found.",
   "ConfigurationProfileBuilder.profileNameRequired": "Profile name is required.",
   "ConfigurationProfileBuilder.canNotUpdateKnownPreset": "Can't add or update a built-in preset configuration profile '%s'.",
-  "AwsClient.configurationProfileNotFound": "Configuration profile '%s' not found."
+  "AwsClient.configurationProfileNotFound": "Configuration profile '%s' not found.",
+  "AwsClient targetClient not defined.": "AwsClient targetClient not defined."
 }

pg/lib/client.ts

@@ -30,6 +30,8 @@ import { RdsMultiAZPgDatabaseDialect } from "./dialect/rds_multi_az_pg_database_
 import { HostInfo } from "../../common/lib/host_info";
 import { TelemetryTraceLevel } from "../../common/lib/utils/telemetry/telemetry_trace_level";
 import { NodePostgresDriverDialect } from "./dialect/node_postgres_driver_dialect";
+import { logTopology } from "../../common/lib/utils/utils";
+import { logger } from "../../common/logutils";
 
 export class AwsPGClient extends AwsClient {
   private static readonly knownDialectsByCode: Map<string, DatabaseDialect> = new Map([

pg/lib/dialect/aurora_pg_database_dialect.ts

@@ -67,7 +67,7 @@ export class AuroraPgDatabaseDialect extends PgDatabaseDialect implements Topolo
 
   async getHostRole(targetClient: ClientWrapper): Promise<HostRole> {
     const res = await targetClient.query(AuroraPgDatabaseDialect.IS_READER_QUERY);
-    return Promise.resolve(res.rows[0]["is_reader"] === "true" ? HostRole.READER : HostRole.WRITER);
+    return Promise.resolve(res.rows[0]["is_reader"] === true ? HostRole.READER : HostRole.WRITER);
   }
 
   async isDialect(targetClient: ClientWrapper): Promise<boolean> {

tests/plugin_manager_benchmarks.ts

@@ -283,7 +283,7 @@ suite(
   }),
 
   add("initHostProviderWith10Plugins", async () => {
-    const pluginManagerWithPlugins = await initPluginManagerWithPlugins(10, instance(mockPluginService), propsWithPlugins);;
+    const pluginManagerWithPlugins = await initPluginManagerWithPlugins(10, instance(mockPluginService), propsWithPlugins);
     return async () =>
       await pluginManagerWithPlugins.initHostProvider(
         new HostInfoBuilder({ hostAvailabilityStrategy: new SimpleHostAvailabilityStrategy() }).withHost("host").build(),
@@ -325,7 +325,7 @@ suite(
   }),
 
   add("notifyConnectionChangedWith10Plugins", async () => {
-    const pluginManagerWithPlugins = await initPluginManagerWithPlugins(10, instance(mockPluginService), propsWithPlugins);;
+    const pluginManagerWithPlugins = await initPluginManagerWithPlugins(10, instance(mockPluginService), propsWithPlugins);
     return async () =>
       await pluginManagerWithPlugins.notifyConnectionChanged(new Set<HostChangeOptions>([HostChangeOptions.INITIAL_CONNECTION]), null);
   }),