query with timeout

aws · Feb 5, 2025 · df73f66 · df73f66
1 parent 7dd4ee2
commit df73f66
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/mysql/lib/dialect/mysql2_driver_dialect.ts b/mysql/lib/dialect/mysql2_driver_dialect.ts
@@ -68,8 +68,8 @@ export class MySQL2DriverDialect implements DriverDialect {
     }
   }
 
-  setQueryTimeout(props: Map<string, any>, sql?: any, wrapperConnectTimeout?: any) {
-    const timeout = wrapperConnectTimeout ?? props.get(WrapperProperties.WRAPPER_QUERY_TIMEOUT.name);
+  setQueryTimeout(props: Map<string, any>, sql?: any, wrapperQueryTimeout?: any) {
+    const timeout = wrapperQueryTimeout ?? props.get(WrapperProperties.WRAPPER_QUERY_TIMEOUT.name);
     if (timeout && !sql[MySQL2DriverDialect.QUERY_TIMEOUT_PROPERTY_NAME]) {
       sql[MySQL2DriverDialect.QUERY_TIMEOUT_PROPERTY_NAME] = Number(timeout);
     }

diff --git a/mysql/lib/dialect/mysql_database_dialect.ts b/mysql/lib/dialect/mysql_database_dialect.ts
@@ -118,14 +118,17 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
 
   async isClientValid(targetClient: ClientWrapper): Promise<boolean> {
     try {
-      return await targetClient
-        .query("SELECT 1")
-        .then(() => {
-          return true;
-        })
-        .catch(() => {
-          return false;
-        });
+      return await ClientUtils.queryWithTimeout(
+        targetClient
+          .query("SELECT 1")
+          .then(() => {
+            return true;
+          })
+          .catch(() => {
+            return false;
+          }),
+        targetClient.properties
+      );
     } catch (error) {
       return false;
     }

diff --git a/pg/lib/dialect/node_postgres_driver_dialect.ts b/pg/lib/dialect/node_postgres_driver_dialect.ts
@@ -71,7 +71,7 @@ export class NodePostgresDriverDialect implements DriverDialect {
     }
   }
 
-  setQueryTimeout(props: Map<string, any>, wrapperQueryTimeout?: any) {
+  setQueryTimeout(props: Map<string, any>, sql?: any, wrapperQueryTimeout?: any) {
     const timeout = wrapperQueryTimeout ?? props.get(WrapperProperties.WRAPPER_QUERY_TIMEOUT.name);
     if (timeout) {
       props.set(NodePostgresDriverDialect.QUERY_TIMEOUT_PROPERTY_NAME, timeout);