chore: refactor wrapper to use the promise API directly (#232)

aws · Oct 9, 2024 · 257df98 · 257df98
1 parent 8323ed3
commit 257df98
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 33 deletions.
diff --git a/common/lib/driver_connection_provider.ts b/common/lib/driver_connection_provider.ts
@@ -53,7 +53,7 @@ export class DriverConnectionProvider implements ConnectionProvider {
     let connectionHostInfo: HostInfo;
 
     try {
-      const targetClient: any = pluginService.createTargetClient(props);
+      const targetClient: any = await Promise.resolve(pluginService.createTargetClient(props));
       const connFunc = pluginService.getConnectFunc(targetClient);
       await connFunc();
       connectionHostInfo = new HostInfoBuilder({

diff --git a/mysql/lib/client.ts b/mysql/lib/client.ts
@@ -17,7 +17,7 @@
 import { QueryOptions } from "mysql2/typings/mysql/lib/protocol/sequences/Query";
 import { AwsClient } from "../../common/lib/aws_client";
 import { WrapperProperties } from "../../common/lib/wrapper_property";
-import { createConnection, Query } from "mysql2";
+import { createConnection, Query } from "mysql2/promise";
 import { MySQLErrorHandler } from "./mysql_error_handler";
 import { MySQLConnectionUrlParser } from "./mysql_connection_url_parser";
 import { DatabaseDialect, DatabaseType } from "../../common/lib/database_dialect/database_dialect";
@@ -66,9 +66,9 @@ export class AwsMySQLClient extends AwsClient {
       this.isConnected = true;
     }
     if (targetClient) {
-      return await ClientUtils.queryWithTimeout(targetClient?.client.promise().query({ sql: sql }), props);
+      return await ClientUtils.queryWithTimeout(targetClient?.client.query({ sql: sql }), props);
     } else {
-      return await ClientUtils.queryWithTimeout(this.targetClient?.client.promise().query({ sql: sql }), props);
+      return await ClientUtils.queryWithTimeout(this.targetClient?.client.query({ sql: sql }), props);
     }
   }
 
@@ -84,7 +84,7 @@ export class AwsMySQLClient extends AwsClient {
       "query",
       async () => {
         await this.pluginService.updateState(options.sql);
-        return await ClientUtils.queryWithTimeout(this.targetClient?.client?.promise().query(options, callback), this.properties);
+        return await ClientUtils.queryWithTimeout(this.targetClient?.client?.query(options, callback), this.properties);
       },
       options
     );
@@ -98,7 +98,7 @@ export class AwsMySQLClient extends AwsClient {
       this.properties,
       "query",
       async () => {
-        return await ClientUtils.queryWithTimeout(this.targetClient?.client?.promise().query(options, callback), this.properties);
+        return await ClientUtils.queryWithTimeout(this.targetClient?.client?.query(options, callback), this.properties);
       },
       options
     );
@@ -209,12 +209,9 @@ export class AwsMySQLClient extends AwsClient {
       "end",
       () => {
         return ClientUtils.queryWithTimeout(
-          this.targetClient?.client
-            ?.promise()
-            .end()
-            .catch((error: any) => {
-              // ignore
-            }),
+          this.targetClient?.client?.end().catch((error: any) => {
+            // ignore
+          }),
           this.properties
         );
       },
@@ -231,7 +228,7 @@ export class AwsMySQLClient extends AwsClient {
       "rollback",
       async () => {
         this.pluginService.updateInTransaction("rollback");
-        return await this.targetClient?.client?.promise().rollback();
+        return await this.targetClient?.client?.rollback();
       },
       null
     );

diff --git a/mysql/lib/dialect/aurora_mysql_database_dialect.ts b/mysql/lib/dialect/aurora_mysql_database_dialect.ts
@@ -40,7 +40,7 @@ export class AuroraMySQLDatabaseDialect extends MySQLDatabaseDialect implements
   }
 
   async queryForTopology(targetClient: ClientWrapper, hostListProvider: HostListProvider): Promise<HostInfo[]> {
-    const res = await targetClient.client.promise().query(AuroraMySQLDatabaseDialect.TOPOLOGY_QUERY);
+    const res = await targetClient.client.query(AuroraMySQLDatabaseDialect.TOPOLOGY_QUERY);
     const hosts: HostInfo[] = [];
     const rows: any[] = res[0];
     rows.forEach((row) => {
@@ -58,18 +58,17 @@ export class AuroraMySQLDatabaseDialect extends MySQLDatabaseDialect implements
   }
 
   async identifyConnection(targetClient: ClientWrapper, props: Map<string, any>): Promise<string> {
-    const res = await targetClient.client.promise().query(AuroraMySQLDatabaseDialect.HOST_ID_QUERY);
+    const res = await targetClient.client.query(AuroraMySQLDatabaseDialect.HOST_ID_QUERY);
     return res[0][0]["host"] ?? "";
   }
 
   async getHostRole(targetClient: ClientWrapper, props: Map<string, any>): Promise<HostRole> {
-    const res = await targetClient.client.promise().query(AuroraMySQLDatabaseDialect.IS_READER_QUERY);
+    const res = await targetClient.client.query(AuroraMySQLDatabaseDialect.IS_READER_QUERY);
     return Promise.resolve(res[0]["is_reader"] === "true" ? HostRole.READER : HostRole.WRITER);
   }
 
   async isDialect(targetClient: ClientWrapper): Promise<boolean> {
     return targetClient.client
-      .promise()
       .query(AuroraMySQLDatabaseDialect.AURORA_VERSION_QUERY)
       .then(([rows]: any) => {
         return !!rows[0]["Value"];

diff --git a/mysql/lib/dialect/mysql_database_dialect.ts b/mysql/lib/dialect/mysql_database_dialect.ts
@@ -42,7 +42,6 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
 
   async getHostAliasAndParseResults(targetClient: ClientWrapper): Promise<string> {
     return targetClient.client
-      .promise()
       .query(this.getHostAliasQuery())
       .then(([rows]: any) => {
         return rows[0]["CONCAT(@@hostname, ':', @@port)"];
@@ -58,7 +57,7 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
 
   async isDialect(targetClient: ClientWrapper): Promise<boolean> {
     return await targetClient.client
-      .promise()
+
       .query(this.getServerVersionQuery())
       .then(([rows]: any) => {
         return rows[0]["Value"].toLowerCase().includes("mysql");
@@ -74,7 +73,7 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
 
   async tryClosingTargetClient(targetClient: ClientWrapper) {
     try {
-      await ClientUtils.queryWithTimeout(targetClient.client.promise().destroy(), targetClient.properties);
+      await ClientUtils.queryWithTimeout(targetClient.client.destroy(), targetClient.properties);
     } catch (error: any) {
       // ignore
     }
@@ -84,7 +83,7 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
     try {
       return await ClientUtils.queryWithTimeout(
         targetClient.client
-          .promise()
+
           .query({ sql: "SELECT 1" })
           .then(() => {
             return true;
@@ -101,12 +100,7 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
 
   getConnectFunc(targetClient: any): () => Promise<any> {
     return async () => {
-      return await targetClient
-        .promise()
-        .connect()
-        .catch((error: any) => {
-          throw error;
-        });
+      return targetClient;
     };
   }
 
@@ -179,6 +173,6 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
   }
 
   async rollback(targetClient: ClientWrapper): Promise<any> {
-    return await targetClient.client.promise().rollback();
+    return await targetClient.client.rollback();
   }
 }
diff --git a/mysql/lib/dialect/rds_multi_az_mysql_database_dialect.ts b/mysql/lib/dialect/rds_multi_az_mysql_database_dialect.ts
@@ -41,7 +41,7 @@ export class RdsMultiAZMySQLDatabaseDialect extends MySQLDatabaseDialect impleme
 
   async isDialect(targetClient: ClientWrapper): Promise<boolean> {
     const res = await targetClient.client
-      .promise()
+
       .query(RdsMultiAZMySQLDatabaseDialect.TOPOLOGY_TABLE_EXIST_QUERY)
       .catch(() => false);
 
@@ -50,7 +50,7 @@ export class RdsMultiAZMySQLDatabaseDialect extends MySQLDatabaseDialect impleme
     }
 
     return !!(await targetClient.client
-      .promise()
+
       .query(RdsMultiAZMySQLDatabaseDialect.TOPOLOGY_QUERY)
       .catch(() => false));
   }
@@ -70,7 +70,7 @@ export class RdsMultiAZMySQLDatabaseDialect extends MySQLDatabaseDialect impleme
         writerHostId = await this.identifyConnection(targetClient, new Map<string, any>());
       }
 
-      const res = await targetClient.client.promise().query(RdsMultiAZMySQLDatabaseDialect.TOPOLOGY_QUERY);
+      const res = await targetClient.client.query(RdsMultiAZMySQLDatabaseDialect.TOPOLOGY_QUERY);
       const rows: any[] = res[0];
       return this.processTopologyQueryResults(hostListProvider, writerHostId, rows);
     } catch (error: any) {
@@ -79,7 +79,7 @@ export class RdsMultiAZMySQLDatabaseDialect extends MySQLDatabaseDialect impleme
   }
 
   private async executeTopologyRelatedQuery(targetClient: ClientWrapper, query: string, resultColumnName?: string): Promise<string> {
-    const res = await targetClient.client.promise().query(query);
+    const res = await targetClient.client.query(query);
     const rows: any[] = res[0];
     if (rows.length > 0) {
       return rows[0][resultColumnName ?? 0];

diff --git a/mysql/lib/dialect/rds_mysql_database_dialect.ts b/mysql/lib/dialect/rds_mysql_database_dialect.ts
@@ -39,7 +39,7 @@ export class RdsMySQLDatabaseDialect extends MySQLDatabaseDialect {
     }
 
     return await targetClient.client
-      .promise()
+
       .query(this.getServerVersionQuery())
       .then(([rows]: any) => {
         return rows[0]["Value"].toLowerCase().includes("source distribution");