query fix

aws · Feb 4, 2025 · 2860b1a · 2860b1a
1 parent 89f2705
commit 2860b1a
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/mysql/lib/client.ts b/mysql/lib/client.ts
@@ -59,7 +59,24 @@ export class AwsMySQLClient extends AwsClient {
   }
 
   async query(options: QueryOptions, callback?: any): Promise<Query> {
-    return await this.targetClient?.client?.query(options, callback);
+    if (!this.isConnected) {
+      await this.connect(); // client.connect is not required for MySQL clients
+      this.isConnected = true;
+    }
+    const host = this.pluginService.getCurrentHostInfo();
+    const context = this.telemetryFactory.openTelemetryContext("awsClient.query", TelemetryTraceLevel.TOP_LEVEL);
+    return await context.start(async () => {
+      return await this.pluginManager.execute(
+        host,
+        this.properties,
+        "query",
+        async () => {
+          await this.pluginService.updateState(options.sql);
+          return await ClientUtils.queryWithTimeout(this.targetClient?.client?.query(options, callback), this.properties);
+        },
+        options
+      );
+    });
   }
 
   private async queryWithoutUpdate(options: QueryOptions, callback?: any): Promise<Query> {