feat: allow custom JSON format in secret value (#602)

Author: karenc-bq

common/lib/authentication/aws_secrets_manager_plugin.ts

@@ -16,6 +16,7 @@
 
 import {
   GetSecretValueCommand,
+  GetSecretValueCommandOutput,
   SecretsManagerClient,
   SecretsManagerClientConfig,
   SecretsManagerServiceException
@@ -39,6 +40,9 @@ export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin implements
   private static SECRETS_ARN_PATTERN: RegExp = new RegExp("^arn:aws:secretsmanager:(?<region>[^:\\n]*):[^:\\n]*:([^:/\\n]*[:/])?(.*)$");
   private readonly pluginService: PluginService;
   private readonly fetchCredentialsCounter;
+  private readonly expirationSec: number;
+  private readonly usernameKey: string;
+  private readonly passwordKey: string;
   private secret: Secret | null = null;
   static secretsCache: Map<string, Secret> = new Map();
   secretKey: SecretCacheKey;
@@ -51,12 +55,25 @@ export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin implements
     const secretId = WrapperProperties.SECRET_ID.get(properties);
     const endpoint = WrapperProperties.SECRET_ENDPOINT.get(properties);
     let region = WrapperProperties.SECRET_REGION.get(properties);
+
+    this.expirationSec = WrapperProperties.SECRET_EXPIRATION_SEC.get(properties);
+    this.usernameKey = WrapperProperties.SECRET_USERNAME_PROPERTY.get(properties);
+    this.passwordKey = WrapperProperties.SECRET_PASSWORD_PROPERTY.get(properties);
+
     const config: SecretsManagerClientConfig = {};
 
     if (!secretId) {
       throw new AwsWrapperError(Messages.get("AwsSecretsManagerConnectionPlugin.missingRequiredConfigParameter", WrapperProperties.SECRET_ID.name));
     }
 
+    if (!this.usernameKey || !this.passwordKey) {
+      throw new AwsWrapperError(Messages.get("AwsSecretsManagerConnectionPlugin.emptyPropertyKeys"));
+    }
+
+    if (this.expirationSec < 0) {
+      throw new AwsWrapperError(Messages.get("AwsSecretsManagerConnectionPlugin.invalidExpirationTime", String(this.expirationSec)));
+    }
+
     if (!region) {
       const groups = secretId.match(AwsSecretsManagerPlugin.SECRETS_ARN_PATTERN)?.groups;
       if (groups?.region) {
@@ -139,12 +156,15 @@ export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin implements
       let fetched = false;
       this.secret = AwsSecretsManagerPlugin.secretsCache.get(JSON.stringify(this.secretKey)) ?? null;
 
-      if (!this.secret || forceRefresh) {
+      if (!this.secret || this.secret.isExpired() || forceRefresh) {
         try {
           this.secret = await this.fetchLatestCredentials();
           fetched = true;
           AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(this.secretKey), this.secret);
         } catch (error: any) {
+          if (error instanceof AwsWrapperError) {
+            throw error;
+          }
           if (error instanceof SecretsManagerServiceException) {
             logAndThrowError(Messages.get("AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials", error.message));
           } else if (error instanceof Error && error.message.includes("AWS SDK error")) {
@@ -163,12 +183,14 @@ export class AwsSecretsManagerPlugin extends AbstractConnectionPlugin implements
       SecretId: this.secretKey.secretId
     };
     const command = new GetSecretValueCommand(commandInput);
-    const result = await this.secretsManagerClient.send(command);
-    const secret = new Secret(JSON.parse(result.SecretString ?? "").username, JSON.parse(result.SecretString ?? "").password);
-    if (secret && secret.username && secret.password) {
-      return secret;
+    const result: GetSecretValueCommandOutput = await this.secretsManagerClient.send(command);
+    const secretJson: string = JSON.parse(result.SecretString ?? "");
+    const username = secretJson[this.usernameKey];
+    const password = secretJson[this.passwordKey];
+    if (!username || !password) {
+      throw new AwsWrapperError(Messages.get("AwsSecretsManagerConnectionPlugin.emptySecretValue", this.usernameKey, this.passwordKey));
     }
-    throw new AwsWrapperError(Messages.get("AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials"));
+    return new Secret(username, password, this.expirationSec);
   }
 
   releaseResources(): Promise<void> {
@@ -198,9 +220,15 @@ export class SecretCacheKey {
 export class Secret {
   readonly username: string;
   readonly password: string;
+  readonly expirationTime: number;
 
-  constructor(username: string, password: string) {
+  constructor(username: string, password: string, expirationSec: number) {
     this.username = username;
     this.password = password;
+    this.expirationTime = Date.now() + expirationSec * 1000;
+  }
+
+  isExpired(): boolean {
+    return Date.now() >= this.expirationTime;
   }
 }

common/lib/utils/messages.ts

@@ -40,7 +40,12 @@ const MESSAGES: Record<string, string> = {
   "HostInfo.weightLessThanZero": "A HostInfo object was created with a weight value less than 0.",
   "AwsSecretsManagerConnectionPlugin.failedToFetchDbCredentials":
     "Was not able to either fetch or read the database credentials from AWS Secrets Manager due to error: %s. Ensure the correct secretId and region properties have been provided.",
+  "AwsSecretsManagerConnectionPlugin.emptySecretValue":
+    "Unable to fetch database credentials with the given username key and password key. Please review the values specified in secretUsernameProperty (%s) and secretPasswordProperty (%s) and ensure they match the Secrets Manager JSON format.",
   "AwsSecretsManagerConnectionPlugin.missingRequiredConfigParameter": "Configuration parameter '%s' is required.",
+  "AwsSecretsManagerConnectionPlugin.emptyPropertyKeys":
+    "secretUsernameProperty and secretPasswordProperty cannot be empty strings. Please ensure they are correct and match the Secret value's JSON format.",
+  "AwsSecretsManagerConnectionPlugin.invalidExpirationTime": "The expiration time (%s) must be set to a non-negative value.",
   "AwsSecretsManagerConnectionPlugin.unhandledError": "Unhandled error: '%s'",
   "AwsSecretsManagerConnectionPlugin.endpointOverrideInvalidConnection": "A connection to the provided endpoint could not be established: '%s'.",
   "ClusterAwareReaderFailoverHandler.invalidTopology": "'%s' was called with an invalid (null or empty) topology",

common/lib/wrapper_property.ts

@@ -168,6 +168,21 @@ export class WrapperProperties {
   static readonly SECRET_ID = new WrapperProperty<string>("secretId", "The name or the ARN of the secret to retrieve.", null);
   static readonly SECRET_REGION = new WrapperProperty<string>("secretRegion", "The region of the secret to retrieve.", null);
   static readonly SECRET_ENDPOINT = new WrapperProperty<string>("secretEndpoint", "The endpoint of the secret to retrieve.", null);
+  static readonly SECRET_EXPIRATION_SEC = new WrapperProperty<number>(
+    "secretExpirationSec",
+    "Secrets Manager credentials' expiration time in seconds.",
+    870
+  );
+  static readonly SECRET_USERNAME_PROPERTY = new WrapperProperty<string>(
+    "secretUsernameProperty",
+    "Set this value to be the key in the JSON secret that contains the username for database connection.",
+    "username"
+  );
+  static readonly SECRET_PASSWORD_PROPERTY = new WrapperProperty<string>(
+    "secretPasswordProperty",
+    "Set this value to be the key in the JSON secret that contains the password for database connection.",
+    "password"
+  );
 
   static readonly FAILOVER_CLUSTER_TOPOLOGY_REFRESH_RATE_MS = new WrapperProperty<number>(
     "failoverClusterTopologyRefreshRateMs",

docs/using-the-nodejs-wrapper/using-plugins/UsingTheAwsSecretsManagerPlugin.md

@@ -20,18 +20,24 @@ The following properties are required for the AWS Secrets Manager Connection Plu
 > [!NOTE]  
 > To use this plugin, you will need to set the following AWS Secrets Manager specific parameters.
 
-| Parameter        | Value  |              Required               | Description                                                                                                                                                                                                                       | Example                  | Default Value |
-| ---------------- | :----: | :---------------------------------: | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :----------------------- | ------------- |
-| `secretId`       | String |                 Yes                 | Set this value to be the secret name or the secret ARN.                                                                                                                                                                           | `secretId`               | `null`        |
-| `secretRegion`   | String | Yes unless the `secretId` is an ARN | Set this value to be the region your secret is in.                                                                                                                                                                                | `us-east-2`              | `null`        |
-| `secretEndpoint` | String |                 No                  | Set this value to be the endpoint override to retrieve your secret from. This parameter value should be in the form of a URL, with a valid protocol (ex. `https://`) and domain (ex. `localhost`). A port number is not required. | `https://localhost:1234` | `null`        |
+| Parameter                |  Value  |              Required               | Description                                                                                                                                                                                                                       | Example                  | Default Value |
+| ------------------------ | :-----: | :---------------------------------: | :-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :----------------------- | ------------- |
+| `secretId`               | String  |                 Yes                 | Set this value to be the secret name or the secret ARN.                                                                                                                                                                           | `secretId`               | `null`        |
+| `secretRegion`           | String  | Yes unless the `secretId` is an ARN | Set this value to be the region your secret is in.                                                                                                                                                                                | `us-east-2`              | `null`        |
+| `secretEndpoint`         | String  |                 No                  | Set this value to be the endpoint override to retrieve your secret from. This parameter value should be in the form of a URL, with a valid protocol (ex. `https://`) and domain (ex. `localhost`). A port number is not required. | `https://localhost:1234` | `null`        |
+| `secretExpirationSec`    | Integer |                 No                  | This property sets the time in seconds that secrets are cached before it is re-fetched.                                                                                                                                           | `600`                    | `870`         |
+| `secretUsernameProperty` | String  |                 No                  | Set this value to be the key in the JSON secret that contains the username for database connection.                                                                                                                               | `db_user`                | `username`    |
+| `secretPasswordProperty` | String  |                 No                  | Set this value to be the key in the JSON secret that contains the password for database connection.                                                                                                                               | `db_pass`                | `password`    |
 
 > [!NOTE]  
 > A Secret ARN has the following format: `arn:aws:secretsmanager:<Region>:<AccountId>:secret:SecretName-6RandomCharacters`
 
 ## Secret Data
 
-The plugin assumes that the secret contains the following properties `username` and `password`.
+The secret stored in the AWS Secrets Manager should be a JSON object containing the properties `username` and `password`. If the secret contains different key names, you can specify them with the `secretUsernameProperty` and `secretPasswordProperty` parameters.
+
+> [!NOTE]
+> Only un-nested JSON format is supported at the moment.
 
 ### Example
 

examples/aws_driver_example/aws_secrets_manager_mysql_example.ts

@@ -20,6 +20,7 @@ const mysqlHost = "db-identifier.XYZ.us-east-2.rds.amazonaws.com";
 const port = 3306;
 const secretId = "SecretName";
 const secretRegion = "us-east-1";
+const secretExpirationTime = 1000;
 /* secretId can be set as secret ARN instead. The ARN includes the secretRegion */
 // const secretId = "arn:aws:secretsmanager:us-east-1:AccountId:secret:SecretName-6RandomCharacters";
 
@@ -29,7 +30,13 @@ const client = new AwsMySQLClient({
   port: port,
   secretId: secretId,
   secretRegion: secretRegion,
-  plugins: "secretsManager"
+  secretExpirationSec: secretExpirationTime,
+  plugins: "secretsManager",
+  // By default, the Secrets Manager plugin assumes the secret stored in the AWS Secrets Manager to be a JSON object containing the properties `username` and `password`.
+  // If the secret contains different key names, you can specify them with the `secretUsernameProperty` and `secretPasswordProperty` parameters.
+  // This example assumes the credentials are stored under the keys db_user and db_pass.
+  secretUsernameProperty: "db_user",
+  secretPasswordProperty: "db_pass"
 });
 
 // Attempt connection.

package.json

@@ -31,7 +31,7 @@
     "mysql"
   ],
   "scripts": {
-    "check": "prettier . --check --config .prettierrc --log-level error",
+    "check": "prettier . --check --config .prettierrc",
     "bench-plugin-manager": "npx tsx tests/plugin_manager_benchmarks.ts",
     "bench-plugins": "npx tsx tests/plugin_benchmarks.ts",
     "bench-plugin-manager-otel": "npx tsx tests/plugin_manager_telemetry_benchmarks.ts",

tests/unit/aws_secrets_manager_plugin.test.ts

@@ -50,7 +50,8 @@ const TEST_HOSTINFO: HostInfo = new HostInfoBuilder({
   hostAvailabilityStrategy: new SimpleHostAvailabilityStrategy(),
   host: TEST_HOST
 }).build();
-const TEST_SECRET = new Secret(TEST_USERNAME, TEST_PASSWORD);
+const EXPIRATION_TIME = 870;
+const TEST_SECRET = new Secret(TEST_USERNAME, TEST_PASSWORD, EXPIRATION_TIME);
 const TEST_SECRET_CACHE_KEY = new SecretCacheKey(TEST_SECRET_ID, TEST_SECRET_REGION);
 
 const VALID_SECRET_RESPONSE = {
@@ -146,7 +147,7 @@ describe("testSecretsManager", () => {
   // In this case, the plugin will fetch the secret and will retry the connection.
   it.each([MYSQL_AUTH_ERROR, PG_AUTH_ERROR])("connect with new secret after trying with cached secrets", async (error) => {
     // Add initial cached secret to be used for a connection.
-    AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(TEST_SECRET_CACHE_KEY), new Secret("", ""));
+    AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(TEST_SECRET_CACHE_KEY), new Secret("", "", EXPIRATION_TIME));
     when(mockSecretsManagerClient.send(anything())).thenResolve(VALID_SECRET_RESPONSE);
     plugin.secretsManagerClient = instance(mockSecretsManagerClient);
 
@@ -210,4 +211,70 @@ describe("testSecretsManager", () => {
     expect(secretKey.region).not.toBe(regionFromArn);
     expect(secretKey.region).toBe(expectedRegion);
   });
+
+  it("connect with custom json keys", async () => {
+    const customSecretResponse = {
+      SecretString: '{"db_user": "foo", "db_pass": "bar"}',
+      $metadata: {}
+    };
+    const props = new Map();
+    WrapperProperties.SECRET_ID.set(props, TEST_SECRET_ID);
+    WrapperProperties.SECRET_REGION.set(props, TEST_SECRET_REGION);
+    WrapperProperties.SECRET_USERNAME_PROPERTY.set(props, "db_user");
+    WrapperProperties.SECRET_PASSWORD_PROPERTY.set(props, "db_pass");
+    const testPlugin = new AwsSecretsManagerPlugin(instance(mockPluginService), props);
+    when(mockSecretsManagerClient.send(anything())).thenResolve(customSecretResponse);
+    testPlugin.secretsManagerClient = instance(mockSecretsManagerClient);
+
+    await testPlugin.connect(TEST_HOSTINFO, props, true, mockConnectFunction);
+
+    verify(mockSecretsManagerClient.send(anything())).once();
+    expect(props.get(WrapperProperties.USER.name)).toBe("foo");
+    expect(props.get(WrapperProperties.PASSWORD.name)).toBe("bar");
+  });
+
+  it("connect with one custom json key", async () => {
+    const customSecretResponse = {
+      SecretString: '{"db_user": "foo", "password": "bar"}',
+      $metadata: {}
+    };
+    const props = new Map();
+    WrapperProperties.SECRET_ID.set(props, TEST_SECRET_ID);
+    WrapperProperties.SECRET_REGION.set(props, TEST_SECRET_REGION);
+    WrapperProperties.SECRET_USERNAME_PROPERTY.set(props, "db_user");
+    const testPlugin = new AwsSecretsManagerPlugin(instance(mockPluginService), props);
+    when(mockSecretsManagerClient.send(anything())).thenResolve(customSecretResponse);
+    testPlugin.secretsManagerClient = instance(mockSecretsManagerClient);
+
+    await testPlugin.connect(TEST_HOSTINFO, props, true, mockConnectFunction);
+
+    verify(mockSecretsManagerClient.send(anything())).once();
+    expect(props.get(WrapperProperties.USER.name)).toBe("foo");
+    expect(props.get(WrapperProperties.PASSWORD.name)).toBe("bar");
+  });
+
+  it("fetch new secret when cached secret is expired", async () => {
+    const expiredSecret = new Secret("oldUser", "oldPass", 0);
+    await new Promise((resolve) => setTimeout(resolve, 10));
+    AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(TEST_SECRET_CACHE_KEY), expiredSecret);
+    when(mockSecretsManagerClient.send(anything())).thenResolve(VALID_SECRET_RESPONSE);
+    plugin.secretsManagerClient = instance(mockSecretsManagerClient);
+
+    await plugin.connect(TEST_HOSTINFO, TEST_PROPS, true, mockConnectFunction);
+
+    verify(mockSecretsManagerClient.send(anything())).once();
+    expect(TEST_PROPS.get(WrapperProperties.USER.name)).toBe(TEST_USERNAME);
+    expect(TEST_PROPS.get(WrapperProperties.PASSWORD.name)).toBe(TEST_PASSWORD);
+  });
+
+  it("use cached secret when not expired", async () => {
+    const validSecret = new Secret(TEST_USERNAME, TEST_PASSWORD, 3600);
+    AwsSecretsManagerPlugin.secretsCache.set(JSON.stringify(TEST_SECRET_CACHE_KEY), validSecret);
+
+    await plugin.connect(TEST_HOSTINFO, TEST_PROPS, true, mockConnectFunction);
+
+    verify(mockSecretsManagerClient.send(anything())).never();
+    expect(TEST_PROPS.get(WrapperProperties.USER.name)).toBe(TEST_USERNAME);
+    expect(TEST_PROPS.get(WrapperProperties.PASSWORD.name)).toBe(TEST_PASSWORD);
+  });
 });