feat: dev plugin (#241)

Co-authored-by: Sophia Chu <[email protected]>

common/lib/aws_client.ts

@@ -157,4 +157,8 @@ export abstract class AwsClient extends EventEmitter {
   }
 
   abstract executeQuery(props: Map<string, any>, query: string, targetClient?: ClientWrapper): any;
+
+  getPluginInstance<T>(iface: any): T {
+    return this.pluginManager.getPluginInstance(iface);
+  }
 }

common/lib/connection_plugin_chain_builder.ts

@@ -34,6 +34,7 @@ import { HostMonitoringPluginFactory } from "./plugins/efm/host_monitoring_plugi
 import { AuroraInitialConnectionStrategyFactory } from "./plugins/aurora_initial_connection_strategy_plugin_factory";
 import { AuroraConnectionTrackerPluginFactory } from "./plugins/connection_tracker/aurora_connection_tracker_plugin_factory";
 import { ConnectionProviderManager } from "./connection_provider_manager";
+import { DeveloperConnectionPluginFactory } from "./plugins/dev/developer_connection_plugin_factory";
 
 /*
   Type alias used for plugin factory sorting. It holds a reference to a plugin
@@ -61,6 +62,7 @@ export class ConnectionPluginChainBuilder {
     ["secretsManager", { factory: AwsSecretsManagerPluginFactory, weight: 1100 }],
     ["federatedAuth", { factory: FederatedAuthPluginFactory, weight: 1200 }],
     ["okta", { factory: OktaAuthPluginFactory, weight: 1300 }],
+    ["dev", { factory: DeveloperConnectionPluginFactory, weight: 1400 }],
     ["connectTime", { factory: ConnectTimePluginFactory, weight: ConnectionPluginChainBuilder.WEIGHT_RELATIVE_TO_PRIOR_PLUGIN }],
     ["executeTime", { factory: ExecuteTimePluginFactory, weight: ConnectionPluginChainBuilder.WEIGHT_RELATIVE_TO_PRIOR_PLUGIN }]
   ]);

common/lib/plugin_manager.ts

@@ -317,4 +317,13 @@ export class PluginManager {
   getTelemetryFactory(): TelemetryFactory {
     return this.telemetryFactory;
   }
+
+  getPluginInstance<T>(iface: any): T {
+    for (const p of this._plugins) {
+      if (p instanceof iface) {
+        return p as T;
+      }
+    }
+    throw new AwsWrapperError(Messages.get("PluginManager.unableToRetrievePlugin"));
+  }
 }

common/lib/plugins/dev/developer_connection_plugin.ts

@@ -0,0 +1,124 @@
+/*
+  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 { PluginService } from "../../plugin_service";
+import { HostInfo } from "../../host_info";
+import { RdsUtils } from "../../utils/rds_utils";
+import { AbstractConnectionPlugin } from "../../abstract_connection_plugin";
+import { ErrorSimulator } from "./error_simulator";
+import { ErrorSimulatorMethodCallback } from "./error_simulator_method_callback";
+import { ErrorSimulatorManager } from "./error_simulator_manager";
+import { logger } from "../../../logutils";
+
+export class DeveloperConnectionPlugin extends AbstractConnectionPlugin implements ErrorSimulator {
+  static ALL_METHODS = "*";
+  static readonly subscribedMethods = new Set<string>(DeveloperConnectionPlugin.ALL_METHODS);
+
+  private errorSimulatorMethodCallback: ErrorSimulatorMethodCallback | null;
+  private nextMethodName: string | null;
+  private nextError: Error | null;
+  pluginService: PluginService;
+  properties: Map<string, any>;
+  rdsUtils: RdsUtils;
+
+  constructor(
+    pluginService: PluginService,
+    properties: Map<string, any>,
+    rdsUtils: RdsUtils = new RdsUtils(),
+    nextMethodName?: string,
+    nextError?: Error,
+    errorSimulatorMethodCallback?: ErrorSimulatorMethodCallback
+  ) {
+    super();
+    this.pluginService = pluginService;
+    this.properties = properties;
+    this.rdsUtils = rdsUtils;
+    this.nextMethodName = nextMethodName ?? null;
+    this.nextError = nextError ?? null;
+    this.errorSimulatorMethodCallback = errorSimulatorMethodCallback ?? null;
+  }
+
+  getSubscribedMethods(): Set<string> {
+    return DeveloperConnectionPlugin.subscribedMethods;
+  }
+
+  raiseErrorOnNextCall(throwable: Error, methodName?: string): void {
+    this.nextError = throwable;
+    this.nextMethodName = methodName ?? DeveloperConnectionPlugin.ALL_METHODS;
+  }
+
+  setCallback(errorSimulatorMethodCallback: ErrorSimulatorMethodCallback): void {
+    this.errorSimulatorMethodCallback = errorSimulatorMethodCallback;
+  }
+
+  override async execute<T>(methodName: string, methodFunc: () => Promise<T>, methodArgs: any[]): Promise<T> {
+    this.raiseErrorIfNeeded(methodName, methodArgs);
+    return methodFunc();
+  }
+
+  raiseErrorIfNeeded<T>(methodName: string, methodArgs: any[]) {
+    if (this.nextError !== null) {
+      if (DeveloperConnectionPlugin.ALL_METHODS === this.nextMethodName || methodName === this.nextMethodName) {
+        this.raiseError(this.nextError, methodName);
+      }
+    } else if (this.errorSimulatorMethodCallback !== null) {
+      this.raiseError(this.errorSimulatorMethodCallback?.getErrorToRaise(methodName, methodArgs), methodName);
+    }
+  }
+
+  raiseError(throwable: Error | null, methodName: string) {
+    if (throwable === null) {
+      return;
+    }
+
+    this.nextError = null;
+    this.nextMethodName = null;
+
+    logger.debug(`Raised an error ${throwable.name} while executing ${methodName}.`);
+
+    throw throwable;
+  }
+
+  connect<T>(hostInfo: HostInfo, props: Map<string, any>, isInitialConnection: boolean, connectFunc: () => Promise<T>): Promise<T> {
+    this.raiseErrorOnConnectIfNeeded(hostInfo, props, isInitialConnection);
+    return connectFunc();
+  }
+
+  forceConnect<T>(hostInfo: HostInfo, props: Map<string, any>, isInitialConnection: boolean, forceConnectFunc: () => Promise<T>): Promise<T> {
+    this.raiseErrorOnConnectIfNeeded(hostInfo, props, isInitialConnection);
+    return forceConnectFunc();
+  }
+
+  raiseErrorOnConnectIfNeeded(hostInfo: HostInfo, props: Map<string, any>, isInitialConnection: boolean) {
+    if (ErrorSimulatorManager.nextError !== null) {
+      this.raiseErrorOnConnect(ErrorSimulatorManager.nextError);
+    } else if (ErrorSimulatorManager.connectCallback !== null) {
+      this.raiseErrorOnConnect(ErrorSimulatorManager.connectCallback.getErrorToRaise(hostInfo, props, isInitialConnection));
+    }
+  }
+
+  raiseErrorOnConnect(throwable: Error | null) {
+    if (!throwable) {
+      return;
+    }
+
+    ErrorSimulatorManager.nextError = null;
+
+    logger.debug(`Raised an error ${throwable.name} while opening a new connection.`);
+
+    throw throwable;
+  }
+}

common/lib/plugins/dev/developer_connection_plugin_factory.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 { ConnectionPluginFactory } from "../../plugin_factory";
+import { PluginService } from "../../plugin_service";
+import { ConnectionPlugin } from "../../connection_plugin";
+import { RdsUtils } from "../../utils/rds_utils";
+import { Messages } from "../../utils/messages";
+import { AwsWrapperError } from "../../utils/errors";
+import { logger } from "../../../logutils";
+
+export class DeveloperConnectionPluginFactory implements ConnectionPluginFactory {
+  async getInstance(pluginService: PluginService, properties: Map<string, any>): Promise<ConnectionPlugin> {
+    try {
+      const developerPlugin = await import("./developer_connection_plugin");
+      return new developerPlugin.DeveloperConnectionPlugin(pluginService, properties, new RdsUtils());
+    } catch (error: any) {
+      logger.error(error);
+      throw new AwsWrapperError(Messages.get("ConnectionPluginChainBuilder.errorImportingPlugin", "DeveloperConnectionPlugin"));
+    }
+  }
+}

common/lib/plugins/dev/error_simulator.ts

@@ -0,0 +1,23 @@
+/*
+  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 { ErrorSimulatorMethodCallback } from "./error_simulator_method_callback";
+
+export interface ErrorSimulator {
+  raiseErrorOnNextCall(throwable: Error, methodName?: string): void;
+
+  setCallback(errorSimulatorMethodCallback: ErrorSimulatorMethodCallback): void;
+}

common/lib/plugins/dev/error_simulator_connect_callback.ts

@@ -0,0 +1,21 @@
+/*
+  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 { HostInfo } from "../../host_info";
+
+export interface ErrorSimulatorConnectCallback {
+  getErrorToRaise(hostInfo: HostInfo, props: Map<string, any>, isInitialConnection: boolean): Error | null;
+}

common/lib/plugins/dev/error_simulator_manager.ts

@@ -0,0 +1,30 @@
+/*
+  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 { ErrorSimulatorConnectCallback } from "./error_simulator_connect_callback";
+
+export class ErrorSimulatorManager {
+  static nextError: Error | null = null;
+  static connectCallback: ErrorSimulatorConnectCallback | null = null;
+
+  static raiseErrorOnNextConnect(throwable: Error): void {
+    ErrorSimulatorManager.nextError = throwable;
+  }
+
+  static setCallback(errorSimulatorConnectCallback: ErrorSimulatorConnectCallback): void {
+    ErrorSimulatorManager.connectCallback = errorSimulatorConnectCallback;
+  }
+}

common/lib/plugins/dev/error_simulator_method_callback.ts

@@ -0,0 +1,19 @@
+/*
+  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.
+*/
+
+export interface ErrorSimulatorMethodCallback {
+  getErrorToRaise<T>(methodName: string, methodArgs: any): Error | null;
+}

common/lib/utils/locales/en.json

@@ -1,6 +1,7 @@
 {
   "PluginManager.unknownPluginCode": "Unknown plugin code: '%s'",
   "PluginManager.PipelineNone": "A pipeline was requested but the created pipeline evaluated to None.",
+  "PluginManager.unableToRetrievePlugin": "Unable to retrieve plugin instance.",
   "ConnectionProvider.unsupportedHostSelectorStrategy": "Unsupported host selection strategy '%s' specified for this connection provider '%s'. Please visit the documentation for all supported strategies.",
   "ConnectionProvider.unsupportedHostInfoSelectorStrategy": "Unsupported host selection strategy '%s' specified for this connection provider '%s'. Please visit the documentation for all supported strategies.",
   "ConnectionPluginChainBuilder.errorImportingPlugin": "The plugin could not be imported. Please ensure the required dependencies have been installed. Plugin: '%s'",

tests/unit/aws_client_get_plugin_instance.test.ts

@@ -0,0 +1,65 @@
+/*
+  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 { PluginService } from "../../common/lib/plugin_service";
+import { instance, mock } from "ts-mockito";
+import { DeveloperConnectionPlugin } from "../../common/lib/plugins/dev/developer_connection_plugin";
+import { RdsUtils } from "../../common/lib/utils/rds_utils";
+import { AwsPGClient } from "../../pg/lib";
+import { IamAuthenticationPlugin } from "../../common/lib/authentication/iam_authentication_plugin";
+
+class DevPluginTest extends DeveloperConnectionPlugin {
+  testMethod() {
+    return "test";
+  }
+}
+
+class TestClient extends AwsPGClient {
+  setManager() {
+    this.pluginManager.init([plugin]);
+  }
+}
+
+const mockPluginService = mock(PluginService);
+const mockRdsUtils = mock(RdsUtils);
+const properties: Map<string, any> = new Map();
+
+const plugin = new DevPluginTest(instance(mockPluginService), properties, instance(mockRdsUtils), undefined, undefined, undefined);
+const testClient = new TestClient({ plugins: "dev" });
+testClient.setManager();
+
+describe("testGetPluginInstance", () => {
+  it("testGetPluginInstanceSameType", async () => {
+    const devPluginTest: DevPluginTest = testClient.getPluginInstance<DevPluginTest>(DeveloperConnectionPlugin);
+    expect(devPluginTest).toEqual(plugin);
+
+    expect(devPluginTest.testMethod()).toEqual("test");
+  });
+
+  it("testGetPluginInstanceAndAssertType", async () => {
+    const developerPlugin: DeveloperConnectionPlugin = testClient.getPluginInstance<DeveloperConnectionPlugin>(DeveloperConnectionPlugin);
+    expect(developerPlugin).toEqual(plugin);
+  });
+
+  it("testGetInstanceWithWrongType", async () => {
+    try {
+      testClient.getPluginInstance(IamAuthenticationPlugin);
+      throw new Error("Retrieved plugin instance of wrong type.");
+    } catch (error: any) {
+      expect(error.message).toEqual("Unable to retrieve plugin instance.");
+    }
+  });
+});