refactor: implement ClientWrapper

Author: karenc-bq

common/lib/aws_pool_client.ts

@@ -17,7 +17,7 @@
 export interface AwsPoolClient {
   connect(): Promise<any>;
 
-  end(poolClient: any): Promise<any>;
+  end(): Promise<any>;
 
   releaseResources(): Promise<void>;
 

common/lib/client_wrapper.ts

@@ -20,4 +20,12 @@ export interface ClientWrapper {
   readonly client: any;
   readonly hostInfo: HostInfo;
   readonly properties: Map<string, any>;
+
+  query(sql: any): Promise<any>;
+
+  end(): Promise<void>;
+
+  rollback(): Promise<void>;
+
+  abort(): Promise<void>;
 }

common/lib/connection_provider.ts

@@ -21,7 +21,6 @@ import { ClientWrapper } from "./client_wrapper";
 
 export interface ConnectionProvider {
   connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ClientWrapper>;
-  end(pluginService: PluginService, clientWrapper: ClientWrapper | undefined): Promise<void>;
   acceptsUrl(hostInfo: HostInfo, props: Map<string, any>): boolean;
   acceptsStrategy(role: HostRole, strategy: string): boolean;
   getHostInfoByStrategy(hosts: HostInfo[], role: HostRole, strategy: string, props?: Map<string, any>): HostInfo;

common/lib/driver_connection_provider.ts

@@ -55,7 +55,7 @@ export class DriverConnectionProvider implements ConnectionProvider {
 
     const driverDialect: DriverDialect = pluginService.getDriverDialect();
     try {
-      const targetClient: any = await driverDialect.connect(props);
+      const targetClient: any = await driverDialect.connect(hostInfo, props);
       connectionHostInfo = new HostInfoBuilder({
         hostAvailabilityStrategy: hostInfo.hostAvailabilityStrategy
       })
@@ -92,12 +92,6 @@ export class DriverConnectionProvider implements ConnectionProvider {
       const originalHost: string = hostInfo.host;
       const fixedHost: string = this.rdsUtils.removeGreenInstancePrefix(hostInfo.host);
       resultProps.set(WrapperProperties.HOST.name, fixedHost);
-      connectionHostInfo = new HostInfoBuilder({
-        hostAvailabilityStrategy: hostInfo.hostAvailabilityStrategy
-      })
-        .copyFrom(hostInfo)
-        .withHost(fixedHost)
-        .build();
 
       logger.info(
         "Connecting to " +
@@ -108,21 +102,10 @@ export class DriverConnectionProvider implements ConnectionProvider {
           JSON.stringify(Object.fromEntries(maskProperties(resultProps)))
       );
 
-      resultTargetClient = driverDialect.connect(resultProps);
+      resultTargetClient = driverDialect.connect(hostInfo, resultProps);
     }
 
-    return {
-      client: resultTargetClient,
-      hostInfo: connectionHostInfo,
-      properties: resultProps
-    };
-  }
-
-  async end(pluginService: PluginService, clientWrapper: ClientWrapper | undefined): Promise<void> {
-    if (clientWrapper === undefined) {
-      return;
-    }
-    return await pluginService.getDriverDialect().end(clientWrapper);
+    return resultTargetClient;
   }
 
   getHostInfoByStrategy(hosts: HostInfo[], role: HostRole, strategy: string, props?: Map<string, any>): HostInfo {

common/lib/driver_dialect/driver_dialect.ts

@@ -17,13 +17,14 @@
 import { ClientWrapper } from "../client_wrapper";
 import { AwsPoolConfig } from "../aws_pool_config";
 import { AwsPoolClient } from "../aws_pool_client";
+import { HostInfo } from "../host_info";
 
 export interface DriverDialect {
   getDialectName(): string;
-  abort(targetClient: ClientWrapper): Promise<void>;
-  rollback(targetClient: ClientWrapper): Promise<any>;
-  connect(props: Map<string, any>): Promise<any>;
-  end(targetClient: ClientWrapper | undefined): Promise<void>;
+
+  connect(hostInfo: HostInfo, props: Map<string, any>): Promise<ClientWrapper>;
+
   preparePoolClientProperties(props: Map<string, any>, poolConfig: AwsPoolConfig | undefined): any;
+
   getAwsPoolClient(props: any): AwsPoolClient;
 }

common/lib/internal_pooled_connection_provider.ts

@@ -37,6 +37,7 @@ import { logger } from "../logutils";
 import { RoundRobinHostSelector } from "./round_robin_host_selector";
 import { AwsPoolClient } from "./aws_pool_client";
 import { AwsPoolConfig } from "./aws_pool_config";
+import { PoolClientWrapper } from "./pool_client_wrapper";
 
 export class InternalPooledConnectionProvider implements PooledConnectionProvider, CanReleaseResources {
   private static readonly acceptedStrategies: Map<string, HostSelector> = new Map([
@@ -108,28 +109,16 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
     const preparedConfig = dialect.preparePoolClientProperties(props, this._poolConfig);
 
     this.internalPool = this.databasePools.computeIfAbsent(
-      new PoolKey(hostInfo.url, this.getPoolKey(hostInfo, props)),
+      new PoolKey(connectionHostInfo.url, this.getPoolKey(connectionHostInfo, props)),
       () => dialect.getAwsPoolClient(preparedConfig),
       this.poolExpirationCheckNanos
     );
 
-    const poolClient = await this.getPoolConnection();
-
-    return {
-      client: poolClient,
-      hostInfo: connectionHostInfo,
-      properties: props
-    };
-  }
-
-  async end(pluginService: PluginService, clientWrapper: ClientWrapper | undefined): Promise<void> {
-    if (this.internalPool && clientWrapper) {
-      return this.internalPool.end(clientWrapper.client);
-    }
+    return await this.getPoolConnection(hostInfo, props);
   }
 
-  async getPoolConnection() {
-    return this.internalPool!.connect();
+  async getPoolConnection(hostInfo: HostInfo, props: Map<string, string>) {
+    return new PoolClientWrapper(await this.internalPool!.connect(), hostInfo, props);
   }
 
   public async releaseResources() {
@@ -172,17 +161,12 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
   }
 
   logConnections() {
-    for (const [key, val] of this.databasePools.entries) {
-      logger.debug(
-        "Internal Pooled Connection: \t\n[ " +
-          key.getPoolKeyString() +
-          "\n\t {\n\t\t\t" +
-          JSON.stringify(val.item.constructor.name) +
-          "\t(expirationTimeNs: " +
-          val.expirationTimeNs +
-          ")\n\t }\n]"
-      );
+    if (this.databasePools.size === 0) {
+      return;
     }
+
+    const l = Array.from(this.databasePools.entries).map(([v, k]) => [v.getPoolKeyString(), k.item.constructor.name]);
+    logger.debug(`Internal Pooled Connection: [${JSON.stringify(l)}]`);
   }
 
   setDatabasePools(connectionPools: SlidingExpirationCache<PoolKey, any>): void {

common/lib/mysql_client_wrapper.ts

@@ -0,0 +1,61 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+import { ClientWrapper } from "./client_wrapper";
+import { HostInfo } from "./host_info";
+import { ClientUtils } from "./utils/client_utils";
+
+/*
+This is an internal wrapper class for the target community driver client created by the MySQL2DriverDialect.
+ */
+export class MySQLClientWrapper implements ClientWrapper {
+  readonly client: any;
+  readonly hostInfo: HostInfo;
+  readonly properties: Map<string, string>;
+
+  /**
+   * Creates a wrapper for the target community driver client.
+   *
+   * @param targetClient The community driver client created for an instance.
+   * @param hostInfo Host information for the connected instance.
+   * @param properties Connection properties for the target client.
+   */
+  constructor(targetClient: any, hostInfo: HostInfo, properties: Map<string, any>) {
+    this.client = targetClient;
+    this.hostInfo = hostInfo;
+    this.properties = properties;
+  }
+
+  query(sql: any): Promise<any> {
+    return this.client?.query(sql);
+  }
+
+  end(): Promise<void> {
+    return this.client?.end();
+  }
+
+  rollback(): Promise<void> {
+    return this.client?.rollback();
+  }
+
+  async abort(): Promise<void> {
+    try {
+      return await ClientUtils.queryWithTimeout(this.client?.destroy(), this.properties);
+    } catch (error: any) {
+      // ignore
+    }
+  }
+}

common/lib/pg_client_wrapper.ts

@@ -0,0 +1,60 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+*/
+
+import { ClientWrapper } from "./client_wrapper";
+import { HostInfo } from "./host_info";
+
+/*
+This an internal wrapper class for a target community driver client created by the NodePostgresPgDriverDialect.
+ */
+export class PgClientWrapper implements ClientWrapper {
+  readonly client: any;
+  readonly hostInfo: HostInfo;
+  readonly properties: Map<string, string>;
+
+  /**
+   * Creates a wrapper for the target community driver client.
+   *
+   * @param targetClient The community driver client created for an instance.
+   * @param hostInfo Host information for the connected instance.
+   * @param properties Connection properties for the target client.
+   */
+  constructor(targetClient: any, hostInfo: HostInfo, properties: Map<string, any>) {
+    this.client = targetClient;
+    this.hostInfo = hostInfo;
+    this.properties = properties;
+  }
+
+  query(sql: any): Promise<any> {
+    return this.client?.query(sql);
+  }
+
+  end(): Promise<void> {
+    return this.client?.end();
+  }
+
+  rollback(): Promise<void> {
+    return this.client?.rollback();
+  }
+
+  async abort(): Promise<void> {
+    try {
+      return await this.end();
+    } catch (error: any) {
+      // Ignore
+    }
+  }
+}

common/lib/plugin_service.ts

@@ -356,7 +356,7 @@ export class PluginService implements ErrorHandler, HostListProviderService {
 
             if (oldClient && (isInTransaction || WrapperProperties.ROLLBACK_ON_SWITCH.get(this.props))) {
               try {
-                await this.getDriverDialect().rollback(oldClient);
+                await oldClient.rollback();
               } catch (error: any) {
                 // Ignore.
               }
@@ -400,13 +400,13 @@ export class PluginService implements ErrorHandler, HostListProviderService {
 
   async abortCurrentClient(): Promise<void> {
     if (this._currentClient.targetClient) {
-      await this.getDriverDialect().abort(this._currentClient.targetClient);
+      await this._currentClient.targetClient.abort();
     }
   }
 
-  async abortTargetClient(targetClient: ClientWrapper | undefined): Promise<void> {
+  async abortTargetClient(targetClient: ClientWrapper | undefined | null): Promise<void> {
     if (targetClient) {
-      await this.getDriverDialect().abort(targetClient);
+      await targetClient.abort();
     }
   }
 
@@ -483,10 +483,6 @@ export class PluginService implements ErrorHandler, HostListProviderService {
     return this._hostListProvider?.getHostRole(client, this.dialect);
   }
 
-  async rollback(targetClient: ClientWrapper) {
-    return await this.getDriverDialect().rollback(targetClient);
-  }
-
   getTelemetryFactory(): TelemetryFactory {
     return this.pluginServiceManagerContainer.pluginManager!.getTelemetryFactory();
   }

common/lib/plugins/failover/writer_failover_handler.ts

@@ -225,9 +225,7 @@ class ReconnectToWriterHandlerTask {
 
     try {
       while (latestTopology.length === 0 && Date.now() < this.endTime && !this.failoverCompleted) {
-        if (this.currentClient) {
-          await this.pluginService.abortTargetClient(this.currentClient);
-        }
+        await this.pluginService.abortTargetClient(this.currentClient);
 
         try {
           const props = new Map(this.initialConnectionProps);
@@ -275,7 +273,7 @@ class ReconnectToWriterHandlerTask {
     this.failoverCompleted = true;
 
     // Task B was returned.
-    if (selectedTask && selectedTask === ClusterAwareWriterFailoverHandler.WAIT_NEW_WRITER_TASK && this.currentClient) {
+    if (selectedTask && selectedTask === ClusterAwareWriterFailoverHandler.WAIT_NEW_WRITER_TASK) {
       await this.pluginService.abortTargetClient(this.currentClient);
     }
   }
@@ -428,7 +426,7 @@ class WaitForNewWriterHandlerTask {
       const props = new Map(this.initialConnectionProps);
       props.set(WrapperProperties.HOST.name, writerCandidate.host);
 
-      let targetClient;
+      let targetClient = null;
       try {
         targetClient = await this.pluginService.forceConnect(writerCandidate, props);
         this.pluginService.setAvailability(writerCandidate.allAliases, HostAvailability.AVAILABLE);
@@ -438,9 +436,7 @@ class WaitForNewWriterHandlerTask {
         return true;
       } catch (error) {
         this.pluginService.setAvailability(writerCandidate.allAliases, HostAvailability.NOT_AVAILABLE);
-        if (targetClient) {
-          await this.pluginService.abortTargetClient(targetClient);
-        }
+        await this.pluginService.abortTargetClient(targetClient);
         return false;
       }
     }
@@ -483,8 +479,6 @@ class WaitForNewWriterHandlerTask {
   }
 
   async callCloseClient(targetClient: ClientWrapper | null) {
-    if (targetClient) {
-      await this.pluginService.abortTargetClient(targetClient);
-    }
+    await this.pluginService.abortTargetClient(targetClient);
   }
 }

common/lib/plugins/stale_dns/stale_dns_helper.ts

@@ -124,7 +124,7 @@ export class StaleDnsHelper {
       logger.debug(Messages.get("StaleDnsHelper.staleDnsDetected", this.writerHostInfo.host));
       this.staleDNSDetectedCounter.inc();
 
-      let targetClient;
+      let targetClient = null;
       try {
         const newProps = new Map<string, any>(props);
         newProps.set(WrapperProperties.HOST.name, this.writerHostInfo.host);