feat: wrapper client setKeepAlive method (#339)

Co-authored-by: Karen Chen <[email protected]>
aws · Dec 11, 2024 · e58e8e9 · e58e8e9
1 parent 02bf0c9
commit e58e8e9
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 16 deletions.
diff --git a/common/lib/database_dialect/database_dialect.ts b/common/lib/database_dialect/database_dialect.ts
@@ -18,8 +18,6 @@ import { HostListProvider } from "../host_list_provider/host_list_provider";
 import { HostListProviderService } from "../host_list_provider_service";
 import { ClientWrapper } from "../client_wrapper";
 import { FailoverRestriction } from "../plugins/failover/failover_restriction";
-import { AwsPoolClient } from "../aws_pool_client";
-import { AwsPoolConfig } from "../aws_pool_config";
 import { ErrorHandler } from "../error_handler";
 
 export enum DatabaseType {

diff --git a/common/lib/driver_dialect/driver_dialect.ts b/common/lib/driver_dialect/driver_dialect.ts
@@ -27,4 +27,6 @@ export interface DriverDialect {
   preparePoolClientProperties(props: Map<string, any>, poolConfig: AwsPoolConfig | undefined): any;
 
   getAwsPoolClient(props: any): AwsPoolClient;
+
+  setKeepAliveProperties(props: Map<string, any>, keepAliveProps: any): void;
 }
diff --git a/common/lib/wrapper_property.ts b/common/lib/wrapper_property.ts
@@ -335,6 +335,12 @@ export class WrapperProperties {
     600_000 // 10 min
   );
 
+  static readonly KEEPALIVE_PROPERTIES = new WrapperProperty<Map<string, any>>(
+    "wrapperKeepAliveProperties",
+    "Map containing any keepAlive properties that the target driver accepts in the client configuration.",
+    null
+  );
+
   static removeWrapperProperties(props: Map<string, any>): any {
     const persistingProperties = [
       WrapperProperties.USER.name,

diff --git a/docs/using-the-nodejs-wrapper/UsingTheNodejsWrapper.md b/docs/using-the-nodejs-wrapper/UsingTheNodejsWrapper.md
diff --git a/mysql/lib/client.ts b/mysql/lib/client.ts
@@ -31,7 +31,6 @@ import { ClientUtils } from "../../common/lib/utils/client_utils";
 import { RdsMultiAZMySQLDatabaseDialect } from "./dialect/rds_multi_az_mysql_database_dialect";
 import { TelemetryTraceLevel } from "../../common/lib/utils/telemetry/telemetry_trace_level";
 import { MySQL2DriverDialect } from "./dialect/mysql2_driver_dialect";
-import { PluginManager } from "../../common/lib";
 
 export class AwsMySQLClient extends AwsClient {
   private static readonly knownDialectsByCode: Map<string, DatabaseDialect> = new Map([

diff --git a/mysql/lib/dialect/mysql2_driver_dialect.ts b/mysql/lib/dialect/mysql2_driver_dialect.ts
@@ -23,6 +23,7 @@ import { AwsPoolClient } from "../../../common/lib/aws_pool_client";
 import { AwsMysqlPoolClient } from "../mysql_pool_client";
 import { MySQLClientWrapper } from "../../../common/lib/mysql_client_wrapper";
 import { HostInfo } from "../../../common/lib/host_info";
+import { UnsupportedMethodError } from "../../../common/lib/utils/errors";
 
 export class MySQL2DriverDialect implements DriverDialect {
   protected dialectName: string = this.constructor.name;
@@ -32,7 +33,10 @@ export class MySQL2DriverDialect implements DriverDialect {
   }
 
   async connect(hostInfo: HostInfo, props: Map<string, any>): Promise<ClientWrapper> {
-    const targetClient = await createConnection(WrapperProperties.removeWrapperProperties(props));
+    const driverProperties = WrapperProperties.removeWrapperProperties(props);
+    // MySQL2 does not support keep alive, explicitly check and throw an exception if this value is set.
+    this.setKeepAliveProperties(driverProperties, props.get(WrapperProperties.KEEPALIVE_PROPERTIES.name));
+    const targetClient = await createConnection(driverProperties);
     return Promise.resolve(new MySQLClientWrapper(targetClient, hostInfo, props));
   }
 
@@ -52,4 +56,10 @@ export class MySQL2DriverDialect implements DriverDialect {
   getAwsPoolClient(props: PoolOptions): AwsPoolClient {
     return new AwsMysqlPoolClient(props);
   }
+
+  setKeepAliveProperties(props: Map<string, any>, keepAliveProps: any) {
+    if (keepAliveProps) {
+      throw new UnsupportedMethodError("Keep alive configuration is not supported for MySQL2.");
+    }
+  }
 }
diff --git a/pg/lib/dialect/node_postgres_driver_dialect.ts b/pg/lib/dialect/node_postgres_driver_dialect.ts
@@ -28,13 +28,17 @@ import { HostInfo } from "../../../common/lib/host_info";
 
 export class NodePostgresDriverDialect implements DriverDialect {
   protected dialectName: string = this.constructor.name;
+  private static keepAlivePropertyName = "keepAlive";
+  private static keepAliveInitialDelayMillisPropertyName = "keepAliveInitialDelayMillis";
 
   getDialectName(): string {
     return this.dialectName;
   }
 
   async connect(hostInfo: HostInfo, props: Map<string, any>): Promise<ClientWrapper> {
-    const targetClient = new pkgPg.Client(WrapperProperties.removeWrapperProperties(props));
+    const driverProperties = WrapperProperties.removeWrapperProperties(props);
+    this.setKeepAliveProperties(driverProperties, props.get(WrapperProperties.KEEPALIVE_PROPERTIES.name));
+    const targetClient = new pkgPg.Client(driverProperties);
     await targetClient.connect();
     return Promise.resolve(new PgClientWrapper(targetClient, hostInfo, props));
   }
@@ -55,4 +59,20 @@ export class NodePostgresDriverDialect implements DriverDialect {
   getAwsPoolClient(props: pkgPg.PoolConfig): AwsPoolClient {
     return new AwsPgPoolClient(props);
   }
+
+  setKeepAliveProperties(props: Map<string, any>, keepAliveProps: any) {
+    if (!keepAliveProps) {
+      return;
+    }
+
+    const keepAlive = keepAliveProps.get(NodePostgresDriverDialect.keepAlivePropertyName);
+    const keepAliveInitialDelayMillis = keepAliveProps.get(NodePostgresDriverDialect.keepAliveInitialDelayMillisPropertyName);
+
+    if (keepAlive) {
+      props.set(NodePostgresDriverDialect.keepAlivePropertyName, keepAlive);
+    }
+    if (keepAliveInitialDelayMillis) {
+      props.set(NodePostgresDriverDialect.keepAliveInitialDelayMillisPropertyName, keepAliveInitialDelayMillis);
+    }
+  }
 }