refactor: implement ClientWrapper

common/lib/aws_pool_client.ts

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

common/lib/client_wrapper.ts

@@ -19,5 +19,9 @@ import { HostInfo } from "./host_info";
 export interface ClientWrapper {
   readonly client: any;
   readonly hostInfo: HostInfo;
-  readonly properties: Map<string, any>;
+  readonly properties: Map<string, string>;
+
+  end(): Promise<void>;
+  rollback(): Promise<void>;
+  destroy(): 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,12 @@
 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(this.internalPool!.connect(), this.internalPool!, hostInfo, props);
   }
 
   public async releaseResources() {

common/lib/mysql_client_wrapper.ts

@@ -0,0 +1,42 @@
+/*
+  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";
+
+export class MySQLClientWrapper implements ClientWrapper {
+  readonly client: any;
+  readonly hostInfo: HostInfo;
+  readonly properties: Map<string, string>;
+
+  constructor(targetClient: any, hostInfo: HostInfo, properties: Map<string, string>) {
+    this.client = targetClient;
+    this.hostInfo = hostInfo;
+    this.properties = properties;
+  }
+
+  end(): Promise<void> {
+    return this.client.end();
+  }
+
+  destroy(): Promise<void> {
+    return this.client.destroy();
+  }
+
+  rollback(): Promise<void> {
+    return this.client.rollback();
+  }
+}

common/lib/pg_client_wrapper.ts

@@ -0,0 +1,42 @@
+/*
+  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";
+
+export class PgClientWrapper implements ClientWrapper {
+  readonly client: any;
+  readonly hostInfo: HostInfo;
+  readonly properties: Map<string, string>;
+
+  constructor(targetClient: any, hostInfo: HostInfo, properties: Map<string, string>) {
+    this.client = targetClient;
+    this.hostInfo = hostInfo;
+    this.properties = properties;
+  }
+
+  end(): Promise<void> {
+    return this.client.end();
+  }
+
+  rollback(): Promise<void> {
+    return this.client.rollback();
+  }
+
+  destroy(): Promise<void> {
+    return this.client.destroy();
+  }
+}

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.
               }
@@ -404,7 +404,7 @@ export class PluginService implements ErrorHandler, HostListProviderService {
     }
   }
 
-  async abortTargetClient(targetClient: ClientWrapper | undefined): Promise<void> {
+  async abortTargetClient(targetClient: ClientWrapper | undefined | null): Promise<void> {
     if (targetClient) {
       await this.getDriverDialect().abort(targetClient);
     }
@@ -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/pool_client_wrapper.ts

@@ -0,0 +1,45 @@
+/*
+  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 { AwsPoolClient } from "./aws_pool_client";
+import { ClientWrapper } from "./client_wrapper";
+import { HostInfo } from "./host_info";
+
+export class PoolClientWrapper implements ClientWrapper {
+  private readonly pool: AwsPoolClient;
+  readonly client: any;
+  readonly hostInfo: HostInfo;
+  readonly properties: Map<string, string>;
+
+  constructor(targetClient: any, pool: AwsPoolClient, hostInfo: HostInfo, properties: Map<string, string>) {
+    this.client = targetClient;
+    this.pool = pool;
+    this.hostInfo = hostInfo;
+    this.properties = properties;
+  }
+
+  end(): Promise<void> {
+    return this.client.release();
+  }
+
+  destroy(): Promise<void> {
+    return this.client.destroy();
+  }
+
+  rollback(): Promise<void> {
+    return this.client.rollback();
+  }
+}

mysql/lib/client.ts

@@ -30,7 +30,6 @@ import { Messages } from "../../common/lib/utils/messages";
 import { ClientWrapper } from "../../common/lib/client_wrapper";
 import { ClientUtils } from "../../common/lib/utils/client_utils";
 import { RdsMultiAZMySQLDatabaseDialect } from "./dialect/rds_multi_az_mysql_database_dialect";
-import { HostInfo } from "../../common/lib/host_info";
 import { TelemetryTraceLevel } from "../../common/lib/utils/telemetry/telemetry_trace_level";
 import { MySQL2DriverDialect } from "./dialect/mysql2_driver_dialect";
 
@@ -220,37 +219,28 @@ export class AwsMySQLClient extends AwsClient {
       return;
     }
 
-    const hostInfo: HostInfo | null = this.pluginService.getCurrentHostInfo();
     const result = await this.pluginManager.execute(
       this.pluginService.getCurrentHostInfo(),
       this.properties,
       "end",
       () => {
-        return ClientUtils.queryWithTimeout(
-          this.pluginService
-            .getConnectionProvider(hostInfo, this.properties)
-            .end(this.pluginService, this.targetClient)
-            .catch((error: any) => {
-              // ignore
-            }),
-          this.properties
-        );
+        return ClientUtils.queryWithTimeout(this.targetClient?.end() ?? Promise.resolve(), this.properties);
       },
       null
     );
     await this.releaseResources();
     return result;
   }
 
-  async rollback(): Promise<Query> {
+  async rollback(): Promise<any> {
     return this.pluginManager.execute(
       this.pluginService.getCurrentHostInfo(),
       this.properties,
       "rollback",
       async () => {
         if (this.targetClient) {
           this.pluginService.updateInTransaction("rollback");
-          return await this.pluginService.getDriverDialect().rollback(this.targetClient);
+          return await this.targetClient.rollback();
         }
         return null;
       },