fix: parameterized queries

.github/workflows/integration_tests_latest.yml

@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - main
+      - fix/parameterized-queries
     paths-ignore:
       - "**/*.md"
       - "**/*.jpg"

common/lib/aws_client.ts

@@ -22,7 +22,6 @@ import { HostListProvider } from "./host_list_provider/host_list_provider";
 import { PluginManager } from "./plugin_manager";
 
 import pkgStream from "stream";
-import { DriverConnectionProvider } from "./driver_connection_provider";
 import { ClientWrapper } from "./client_wrapper";
 import { ConnectionProviderManager } from "./connection_provider_manager";
 import { DefaultTelemetryFactory } from "./utils/telemetry/default_telemetry_factory";
@@ -35,10 +34,12 @@ import { AwsWrapperError } from "./utils/errors";
 import { Messages } from "./utils/messages";
 import { TransactionIsolationLevel } from "./utils/transaction_isolation_level";
 import { HostListProviderService } from "./host_list_provider_service";
+import { SessionStateClient } from "./session_state_client";
+import { ConnectionProvider } from "./connection_provider";
 
 const { EventEmitter } = pkgStream;
 
-export abstract class AwsClient extends EventEmitter {
+export abstract class AwsClient extends EventEmitter implements SessionStateClient {
   private _defaultPort: number = -1;
   protected telemetryFactory: TelemetryFactory;
   protected pluginManager: PluginManager;
@@ -55,7 +56,8 @@ export abstract class AwsClient extends EventEmitter {
     dbType: DatabaseType,
     knownDialectsByCode: Map<string, DatabaseDialect>,
     parser: ConnectionUrlParser,
-    driverDialect: DriverDialect
+    driverDialect: DriverDialect,
+    connectionProvider: ConnectionProvider
   ) {
     super();
     this.config = config;
@@ -110,7 +112,7 @@ export abstract class AwsClient extends EventEmitter {
     this.pluginManager = new PluginManager(
       container,
       this.properties,
-      new ConnectionProviderManager(new DriverConnectionProvider(), WrapperProperties.CONNECTION_PROVIDER.get(this.properties)),
+      new ConnectionProviderManager(connectionProvider, WrapperProperties.CONNECTION_PROVIDER.get(this.properties)),
       this.telemetryFactory
     );
   }

common/lib/client_wrapper.ts

@@ -22,7 +22,9 @@ export interface ClientWrapper {
   readonly properties: Map<string, any>;
   readonly id: string;
 
+  // Internal method, executes wrapper-specific queries like the topology query.
   query(sql: string): Promise<any>;
+  query(config: any, values?: any, callback?: any): Promise<any>;
 
   end(): Promise<void>;
 

common/lib/driver_dialect/driver_dialect.ts

@@ -33,4 +33,6 @@ export interface DriverDialect {
   setQueryTimeout(props: Map<string, any>, sql?: any, wrapperConnectTimeout?: any): void;
 
   setKeepAliveProperties(props: Map<string, any>, keepAliveProps: any): void;
+
+  getQueryFromMethodArg(methodArg: any): string;
 }

common/lib/host_list_provider_service.ts

@@ -21,7 +21,7 @@ import { DatabaseDialect } from "./database_dialect/database_dialect";
 import { HostInfoBuilder } from "./host_info_builder";
 import { ConnectionUrlParser } from "./utils/connection_url_parser";
 import { TelemetryFactory } from "./utils/telemetry/telemetry_factory";
-import { AllowedAndBlockedHosts } from "./AllowedAndBlockedHosts";
+import { AllowedAndBlockedHosts } from "./allowed_and_blocked_hosts";
 
 export interface HostListProviderService {
   getHostListProvider(): HostListProvider | null;

common/lib/pg_client_wrapper.ts

@@ -41,8 +41,8 @@ export class PgClientWrapper implements ClientWrapper {
     this.id = uniqueId("PgClient_");
   }
 
-  query(sql: string): Promise<any> {
-    return this.client?.query(sql);
+  query(config: any, values?: any, callback?: any): Promise<any> {
+    return this.client?.query(config, values, callback);
   }
 
   queryWithTimeout(sql: string): Promise<any> {

common/lib/plugin_service.ts

@@ -43,7 +43,7 @@ import { DatabaseDialectCodes } from "./database_dialect/database_dialect_codes"
 import { getWriter, logTopology } from "./utils/utils";
 import { TelemetryFactory } from "./utils/telemetry/telemetry_factory";
 import { DriverDialect } from "./driver_dialect/driver_dialect";
-import { AllowedAndBlockedHosts } from "./AllowedAndBlockedHosts";
+import { AllowedAndBlockedHosts } from "./allowed_and_blocked_hosts";
 import { ConnectionPlugin } from "./connection_plugin";
 
 export interface PluginService extends ErrorHandler {

common/lib/plugins/custom_endpoint/custom_endpoint_monitor_impl.ts

@@ -20,7 +20,7 @@ import { PluginService } from "../../plugin_service";
 import { TelemetryCounter } from "../../utils/telemetry/telemetry_counter";
 import { logger } from "../../../logutils";
 import { CustomEndpointInfo } from "./custom_endpoint_info";
-import { AllowedAndBlockedHosts } from "../../AllowedAndBlockedHosts";
+import { AllowedAndBlockedHosts } from "../../allowed_and_blocked_hosts";
 import { CacheMap } from "../../utils/cache_map";
 import { MemberListType } from "./member_list_type";
 import { Messages } from "../../utils/messages";

common/lib/plugins/read_write_splitting_plugin.ts

@@ -142,7 +142,7 @@ export class ReadWriteSplittingPlugin extends AbstractConnectionPlugin implement
   }
 
   override async execute<T>(methodName: string, executeFunc: () => Promise<T>, methodArgs: any): Promise<T> {
-    const statement = methodArgs.sql ?? methodArgs;
+    const statement = SqlMethodUtils.parseMethodArgs(methodArgs, this.pluginService.getDriverDialect());
     const statements = SqlMethodUtils.parseMultiStatementQueries(statement);
 
     const updateReadOnly: boolean | undefined = SqlMethodUtils.doesSetReadOnly(statements, this.pluginService.getDialect());

common/lib/session_state_client.ts

@@ -0,0 +1,39 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { TransactionIsolationLevel } from "./utils/transaction_isolation_level";
+
+export interface SessionStateClient {
+  setReadOnly(readOnly: boolean): Promise<any | void>;
+
+  isReadOnly(): boolean;
+
+  setAutoCommit(autoCommit: boolean): Promise<any | void>;
+
+  getAutoCommit(): boolean;
+
+  setTransactionIsolation(level: TransactionIsolationLevel): Promise<any | void>;
+
+  getTransactionIsolation(): TransactionIsolationLevel;
+
+  setSchema(schema: any): Promise<any | void>;
+
+  getSchema(): string;
+
+  setCatalog(catalog: string): Promise<any | void>;
+
+  getCatalog(): string;
+}

common/lib/utils/sql_method_utils.ts

@@ -16,6 +16,7 @@
 
 import { DatabaseDialect } from "../database_dialect/database_dialect";
 import { TransactionIsolationLevel } from "./transaction_isolation_level";
+import { DriverDialect } from "../driver_dialect/driver_dialect";
 
 export class SqlMethodUtils {
   static doesOpenTransaction(sql: string) {
@@ -129,4 +130,17 @@ export class SqlMethodUtils {
 
     return sql.split(";");
   }
+
+  static parseMethodArgs(methodArgs: any, driverDialect: DriverDialect) {
+    // MethodArgs may be an array, where the first element could be either a string, a MySQL2 Query Object, or a Node-Postgres config object.
+    if (!methodArgs) {
+      return methodArgs;
+    }
+    if (!Array.isArray(methodArgs)) {
+      return driverDialect.getQueryFromMethodArg(methodArgs);
+    }
+
+    const statement = methodArgs[0];
+    return driverDialect.getQueryFromMethodArg(statement);
+  }
 }

common/lib/utils/subscribed_method_helper.ts

@@ -15,6 +15,35 @@
 */
 
 export class SubscribedMethodHelper {
-  static readonly NETWORK_BOUND_METHODS: string[] = ["connect", "forceConnect", "query", "rollback"];
-  static readonly METHODS_REQUIRING_UPDATED_TOPOLOGY: string[] = ["connect", "forceConnect", "query"];
+  static readonly NETWORK_BOUND_METHODS: string[] = [
+    "connect",
+    "forceConnect",
+    "query",
+    // MySQL-specific
+    "execute",
+    "rollback",
+    "beginTransaction",
+    "commit",
+    "changeUser",
+    "pause",
+    "resume",
+    "prepare",
+    "unprepare"
+  ];
+
+  static readonly METHODS_REQUIRING_UPDATED_TOPOLOGY: string[] = [
+    "connect",
+    "forceConnect",
+    "query",
+    // MySQL-specific
+    "beginTransaction",
+    "commit",
+    "changeUser",
+    "pause",
+    "resume",
+    "execute",
+    "rollback",
+    "prepare",
+    "unprepare"
+  ];
 }

common/lib/wrapper_property.ts

@@ -464,14 +464,8 @@ export class WrapperProperties {
     BlueGreenStatusProvider.MONITORING_PROPERTY_PREFIX
   ];
 
-  private static startsWithPrefix(key: string) {
-    for (const prefix in WrapperProperties.PREFIXES) {
-      if (key.startsWith(prefix)) {
-        return true;
-      }
-    }
-
-    return false;
+  private static startsWithPrefix(key: string): boolean {
+    return WrapperProperties.PREFIXES.some((prefix) => key.startsWith(prefix));
   }
 
   static removeWrapperProperties(props: Map<string, any>): Map<string, any> {