chore: refactor read write splitting

Author: karenc-bq

.github/workflows/main.yml

@@ -8,6 +8,7 @@ on:
   pull_request:
     branches:
       - "*"
+      - dev/v3
 
 permissions:
   contents: read

common/lib/connection_info.ts

@@ -0,0 +1,35 @@
+/*
+ * 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 { ClientWrapper } from "./client_wrapper";
+
+export class ConnectionInfo {
+  private readonly _client: ClientWrapper;
+  private readonly _isPooled: boolean;
+
+  constructor(client: ClientWrapper, isPooled: boolean) {
+    this._client = client;
+    this._isPooled = isPooled;
+  }
+
+  get client(): ClientWrapper {
+    return this._client;
+  }
+
+  get isPooled(): boolean {
+    return this._isPooled;
+  }
+}

common/lib/connection_provider.ts

@@ -17,10 +17,10 @@
 import { HostRole } from "./host_role";
 import { HostInfo } from "./host_info";
 import { PluginService } from "./plugin_service";
-import { ClientWrapper } from "./client_wrapper";
+import { ConnectionInfo } from "./connection_info";
 
 export interface ConnectionProvider {
-  connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ClientWrapper>;
+  connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ConnectionInfo>;
   acceptsUrl(hostInfo: HostInfo, props: Map<string, any>): boolean;
   acceptsStrategy(role: HostRole, strategy: string): boolean;
   getHostInfoByStrategy(hosts: HostInfo[], role: HostRole, strategy: string, props?: Map<string, any>): HostInfo;

common/lib/driver_connection_provider.ts

@@ -30,6 +30,7 @@ import { logger } from "../logutils";
 import { ClientWrapper } from "./client_wrapper";
 import { RoundRobinHostSelector } from "./round_robin_host_selector";
 import { DriverDialect } from "./driver_dialect/driver_dialect";
+import { ConnectionInfo } from "./connection_info";
 
 export class DriverConnectionProvider implements ConnectionProvider {
   private static readonly acceptedStrategies: Map<string, HostSelector> = new Map([
@@ -46,7 +47,7 @@ export class DriverConnectionProvider implements ConnectionProvider {
     return true;
   }
 
-  async connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ClientWrapper> {
+  async connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ConnectionInfo> {
     let resultTargetClient;
     const resultProps = new Map(props);
     resultProps.set(WrapperProperties.HOST.name, hostInfo.host);
@@ -92,7 +93,7 @@ export class DriverConnectionProvider implements ConnectionProvider {
       resultTargetClient = driverDialect.connect(hostInfo, resultProps);
     }
     pluginService.attachErrorListener(resultTargetClient);
-    return resultTargetClient;
+    return new ConnectionInfo(resultTargetClient, false);
   }
 
   getHostInfoByStrategy(hosts: HostInfo[], role: HostRole, strategy: string, props?: Map<string, any>): HostInfo {

common/lib/internal_pooled_connection_provider.ts

@@ -39,6 +39,7 @@ import { LeastConnectionsHostSelector } from "./least_connections_host_selector"
 import { PoolClientWrapper } from "./pool_client_wrapper";
 import { logger } from "../logutils";
 import { SlidingExpirationCacheWithCleanupTask } from "./utils/sliding_expiration_cache_with_cleanup_task";
+import { ConnectionInfo } from "./connection_info";
 
 export class InternalPooledConnectionProvider implements PooledConnectionProvider, CanReleaseResources {
   static readonly CACHE_CLEANUP_NANOS: bigint = BigInt(10 * 60_000_000_000); // 10 minutes
@@ -79,7 +80,7 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
     return RdsUrlType.RDS_INSTANCE === urlType;
   }
 
-  async connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ClientWrapper> {
+  async connect(hostInfo: HostInfo, pluginService: PluginService, props: Map<string, any>): Promise<ConnectionInfo> {
     const resultProps = new Map(props);
     resultProps.set(WrapperProperties.HOST.name, hostInfo.host);
     if (hostInfo.isPortSpecified()) {
@@ -122,7 +123,7 @@ export class InternalPooledConnectionProvider implements PooledConnectionProvide
 
     const poolClient = await this.getPoolConnection(connectionHostInfo, props);
     pluginService.attachErrorListener(poolClient);
-    return poolClient;
+    return new ConnectionInfo(poolClient, true);
   }
 
   async getPoolConnection(hostInfo: HostInfo, props: Map<string, string>) {

common/lib/plugin_service.ts

@@ -150,6 +150,10 @@ export interface PluginService extends ErrorHandler {
   getStatus<T>(clazz: any, key: string): T;
 
   isPluginInUse(plugin: any): boolean;
+
+  isPooledClient(): boolean;
+
+  setIsPooledClient(isPooledClient: boolean): void;
 }
 
 export class PluginServiceImpl implements PluginService, HostListProviderService {
@@ -172,6 +176,8 @@ export class PluginServiceImpl implements PluginService, HostListProviderService
   protected static readonly statusesExpiringCache: CacheMap<string, any> = new CacheMap();
   protected static readonly DEFAULT_STATUS_CACHE_EXPIRE_NANO: number = 3_600_000_000_000; // 60 minutes
 
+  protected pooledClient: boolean | null = null;
+
   constructor(
     container: PluginServiceManagerContainer,
     client: AwsClient,
@@ -782,4 +788,12 @@ export class PluginServiceImpl implements PluginService, HostListProviderService
   isPluginInUse(plugin: any) {
     return this.pluginServiceManagerContainer.pluginManager!.isPluginInUse(plugin);
   }
+
+  isPooledClient(): boolean | null {
+    return this.pooledClient;
+  }
+
+  setIsPooledClient(isPooledClient: boolean): void {
+    this.pooledClient = isPooledClient;
+  }
 }

common/lib/plugins/default_plugin.ts

@@ -29,6 +29,7 @@ import { AwsWrapperError } from "../utils/errors";
 import { HostAvailability } from "../host_availability/host_availability";
 import { ClientWrapper } from "../client_wrapper";
 import { TelemetryTraceLevel } from "../utils/telemetry/telemetry_trace_level";
+import { ConnectionInfo } from "../connection_info";
 
 export class DefaultPlugin extends AbstractConnectionPlugin {
   id: string = uniqueId("_defaultPlugin");
@@ -79,10 +80,11 @@ export class DefaultPlugin extends AbstractConnectionPlugin {
       TelemetryTraceLevel.NESTED
     );
 
-    const result = await telemetryContext.start(async () => await connProvider.connect(hostInfo, this.pluginService, props));
+    const result: ConnectionInfo = await telemetryContext.start(async () => await connProvider.connect(hostInfo, this.pluginService, props));
     this.pluginService.setAvailability(hostInfo.allAliases, HostAvailability.AVAILABLE);
-    await this.pluginService.updateDialect(result);
-    return result;
+    this.pluginService.setIsPooledClient(result.isPooled);
+    await this.pluginService.updateDialect(result.client);
+    return result.client;
   }
 
   override async execute<Type>(methodName: string, methodFunc: () => Promise<Type>): Promise<Type> {