diff --git a/common/lib/aws_client.ts b/common/lib/aws_client.ts index 0ef229a7..39c08822 100644 --- a/common/lib/aws_client.ts +++ b/common/lib/aws_client.ts @@ -145,8 +145,6 @@ export abstract class AwsClient extends EventEmitter { await this.pluginManager.releaseResources(); } - abstract executeQuery(props: Map, query: string, targetClient?: ClientWrapper): any; - getPluginInstance(iface: any): T { return this.pluginManager.getPluginInstance(iface); } diff --git a/common/lib/client_wrapper.ts b/common/lib/client_wrapper.ts index b16b4b30..68ac21df 100644 --- a/common/lib/client_wrapper.ts +++ b/common/lib/client_wrapper.ts @@ -29,4 +29,6 @@ export interface ClientWrapper { rollback(): Promise; abort(): Promise; + + queryWithTimeout(sql: string): Promise; } diff --git a/common/lib/mysql_client_wrapper.ts b/common/lib/mysql_client_wrapper.ts index 5810c985..d79d3907 100644 --- a/common/lib/mysql_client_wrapper.ts +++ b/common/lib/mysql_client_wrapper.ts @@ -46,6 +46,10 @@ export class MySQLClientWrapper implements ClientWrapper { return this.client?.query(sql); } + async queryWithTimeout(sql: string): Promise { + return await ClientUtils.queryWithTimeout(this.client.query({ sql: sql }), this.properties); + } + end(): Promise { return this.client?.end(); } diff --git a/common/lib/pg_client_wrapper.ts b/common/lib/pg_client_wrapper.ts index 645d4bc8..8989315e 100644 --- a/common/lib/pg_client_wrapper.ts +++ b/common/lib/pg_client_wrapper.ts @@ -45,6 +45,10 @@ export class PgClientWrapper implements ClientWrapper { return this.client?.query(sql); } + queryWithTimeout(sql: string): Promise { + return this.client?.queryWithTimeout(this.client.query(sql), this.properties); + } + end(): Promise { return this.client?.end(); } diff --git a/common/lib/pool_client_wrapper.ts b/common/lib/pool_client_wrapper.ts index 3cbe8db3..0cdacb67 100644 --- a/common/lib/pool_client_wrapper.ts +++ b/common/lib/pool_client_wrapper.ts @@ -17,6 +17,7 @@ import { ClientWrapper } from "./client_wrapper"; import { HostInfo } from "./host_info"; import { uniqueId } from "../logutils"; +import { ClientUtils } from "./utils/client_utils"; export class PoolClientWrapper implements ClientWrapper { readonly client: any; @@ -39,6 +40,10 @@ export class PoolClientWrapper implements ClientWrapper { return this.client?.query(sql); } + async queryWithTimeout(sql: string): Promise { + return await ClientUtils.queryWithTimeout(this.client.query(sql), this.properties); + } + async end(): Promise { try { return this.client?.release(); diff --git a/mysql/lib/client.ts b/mysql/lib/client.ts index a3a31dd6..2e0e17e3 100644 --- a/mysql/lib/client.ts +++ b/mysql/lib/client.ts @@ -61,18 +61,6 @@ export class AwsMySQLClient extends AwsClient { }); } - async executeQuery(props: Map, sql: string, targetClient?: ClientWrapper): Promise { - if (!this.isConnected) { - await this.connect(); // client.connect is not required for MySQL clients - this.isConnected = true; - } - if (targetClient) { - return await ClientUtils.queryWithTimeout(targetClient?.client.query({ sql: sql }), props); - } else { - return await ClientUtils.queryWithTimeout(this.targetClient?.client.query({ sql: sql }), props); - } - } - async query(options: QueryOptions, callback?: any): Promise { if (!this.isConnected) { await this.connect(); // client.connect is not required for MySQL clients @@ -109,7 +97,7 @@ export class AwsMySQLClient extends AwsClient { } async updateSessionStateReadOnly(readOnly: boolean): Promise { - const result = await this.executeQuery(this.properties, `SET SESSION TRANSACTION READ ${readOnly ? "ONLY" : "WRITE"}`, this.targetClient); + const result = await this.targetClient.queryWithTimeout(`SET SESSION TRANSACTION READ ${readOnly ? "ONLY" : "WRITE"}`); this._isReadOnly = readOnly; this.pluginService.getSessionStateService().setupPristineReadOnly(); diff --git a/pg/lib/client.ts b/pg/lib/client.ts index c57ee05a..a836e9b9 100644 --- a/pg/lib/client.ts +++ b/pg/lib/client.ts @@ -59,14 +59,6 @@ export class AwsPGClient extends AwsClient { }); } - executeQuery(props: Map, sql: string, targetClient?: ClientWrapper): Promise { - if (targetClient) { - return targetClient?.client.query(sql); - } else { - return this.targetClient?.client.query(sql); - } - } - async query(text: string): Promise { const context = this.telemetryFactory.openTelemetryContext("awsClient.query", TelemetryTraceLevel.TOP_LEVEL); return await context.start(async () => { @@ -84,11 +76,7 @@ export class AwsPGClient extends AwsClient { } async updateSessionStateReadOnly(readOnly: boolean): Promise { - if (readOnly) { - return await this.executeQuery(this.properties, "SET SESSION CHARACTERISTICS AS TRANSACTION READ ONLY"); - } else { - return await this.executeQuery(this.properties, "SET SESSION CHARACTERISTICS AS TRANSACTION READ WRITE"); - } + return await this.targetClient.query(`SET SESSION CHARACTERISTICS AS TRANSACTION READ ${readOnly ? "ONLY" : "WRITE"}`); } private async readOnlyQuery(text: string): Promise {