feat(Model): fetch tableName from (SSM) Parameter Store (#291)

* feat(Model): fetch tableName from (SSM) Parameter Store

* fix: fix tests & rewrite SSMParam decorator
* fix: add missing dependency
* docs(readme): add SSM parameter feature
* fix(batchWrite): await tableName

---------

Co-authored-by: ThibaultRby <[email protected]>

Author: SajidHamza9

README.md

@@ -61,6 +61,29 @@ export class Album extends Model<IAlbum> {
 }
 ```
 
+You have the option to either directly provide the table name, as shown in the previous example, or retrieve it from the AWS Systems Manager (SSM) Parameter Store by specifying the ARN of the SSM parameter. The value will be automatically replaced at runtime.
+
+Here is an example of using table name as SSM Parameter ARN
+
+```typescript
+interface IAlbum {
+  artist: string;
+  album: string;
+  year?: number;
+  genres?: string[];
+}
+
+export class Album extends Model<IAlbum> {
+  protected tableName = 'arn:aws:ssm:<aws-region>:<aws-account-id>:parameter/ParameterName';
+  protected hashkey = 'artist';
+  protected rangekey = 'album';
+
+  constructor(item?: IAlbum) {
+    super(item);
+  }
+}
+```
+
 Here is another example for a table with a simple hashkey:
 
 ```typescript

package-lock.json

@@ -19,8 +19,10 @@
   },
   "dependencies": {
     "@aws-sdk/client-dynamodb": "^3.445.0",
+    "@aws-sdk/client-ssm": "^3.496.0",
     "@aws-sdk/lib-dynamodb": "^3.445.0",
     "joi": "^17.11.0",
+    "reflect-metadata": "^0.2.1",
     "uuid": "^9.0.1"
   },
   "devDependencies": {

package.json

@@ -25,6 +25,7 @@ import Scan from './scan';
 import { IUpdateActions, buildUpdateActions, put } from './update-operators';
 import ValidationError from './validation-error';
 import { PutCommandInput } from '@aws-sdk/lib-dynamodb/dist-types/commands/PutCommand';
+import { GetParameterCommand, SSMClient } from '@aws-sdk/client-ssm';
 
 export type KeyValue = string | number | Buffer | boolean | null;
 type SimpleKey = KeyValue;
@@ -36,8 +37,45 @@ const isCompositeKey = (hashKeys_compositeKeys: Keys): hashKeys_compositeKeys is
 
 const isSimpleKey = (hashKeys_compositeKeys: Keys): hashKeys_compositeKeys is SimpleKey[] => hashKeys_compositeKeys.length > 0 && Model.isKey(hashKeys_compositeKeys[0]);
 
+export const fetchTableName = async (tableName: string) => {
+  if (tableName.startsWith('arn:aws:ssm')) {
+    const ssmArnRegex = /^arn:aws:ssm:([a-z0-9-]+):\d{12}:parameter\/([a-zA-Z0-9_.\-/]+)$/;
+    const isMatchingArn = tableName.match(ssmArnRegex)
+    if (!isMatchingArn) {
+      throw new Error("Invalid syntax for table name as SSM Parameter");
+    }
+    const [_, region, parameterName] = isMatchingArn;
+    const ssmClient = new SSMClient({ region });
+    const getValue = new GetParameterCommand({
+      Name: parameterName,
+    });
+    try {
+      const { Parameter } = await ssmClient.send(getValue);
+      return Parameter?.Value;
+    }
+    catch (e) {
+      throw new Error("Invalid SSM Parameter");
+    }
+  }
+  return tableName;
+}
+
+const SSMParam = (target: any, key: string) => {
+  const symbol = Symbol();
+  Reflect.defineProperty(target, key, {
+    get: function () {
+      return (async () => await this[symbol])();
+    },
+    set: function (newVal: string) {
+      this[symbol] = fetchTableName(newVal);
+    }
+  })
+}
+
 export default abstract class Model<T> {
-  protected tableName: string | undefined;
+
+  @SSMParam
+  protected tableName: string | Promise<string | undefined> | undefined;
 
   protected item: T | undefined;
 
@@ -235,7 +273,7 @@ export default abstract class Model<T> {
     }
     // Prepare putItem operation
     const params: PutCommandInput = {
-      TableName: this.tableName,
+      TableName: await this.tableName,
       Item: toSave,
     };
     // Overload putItem parameters with options given in arguments (if any)
@@ -275,7 +313,7 @@ export default abstract class Model<T> {
     // Prepare getItem operation
     this.testKeys(pk, sk);
     const params: GetCommandInput = {
-      TableName: this.tableName,
+      TableName: await this.tableName,
       Key: this.buildKeys(pk, sk),
     };
     // Overload getItem parameters with options given in arguments (if any)
@@ -376,7 +414,7 @@ export default abstract class Model<T> {
       throw new Error('Item to delete does not exists');
     }
     const params: DeleteCommandInput = {
-      TableName: this.tableName,
+      TableName: await this.tableName,
       Key: this.buildKeys(pk, sk),
     };
     if (options) {
@@ -399,12 +437,12 @@ export default abstract class Model<T> {
       throw new Error('Primary key is not defined on your model');
     }
     const params: ScanCommandInput = {
-      TableName: this.tableName,
+      TableName: '',
     };
     if (options) {
       Object.assign(params, options);
     }
-    return new Scan(this.documentClient, params, this.pk, this.sk);
+    return new Scan(this.documentClient, params, this.tableName, this.pk, this.sk);
   }
 
   /**
@@ -445,15 +483,15 @@ export default abstract class Model<T> {
         : (index_options as Partial<QueryCommandInput>);
     // Building query
     const params: QueryCommandInput = {
-      TableName: this.tableName,
+      TableName: '',
     };
     if (indexName) {
       params.IndexName = indexName;
     }
     if (queryOptions) {
       Object.assign(params, queryOptions);
     }
-    return new Query(this.documentClient, params, this.pk, this.sk);
+    return new Query(this.documentClient, params, this.tableName, this.pk, this.sk);
   }
 
   /**
@@ -473,15 +511,15 @@ export default abstract class Model<T> {
     if (isCompositeKey(keys)) {
       params = {
         RequestItems: {
-          [this.tableName]: {
+          [await this.tableName as string]: {
             Keys: keys.map((k) => this.buildKeys(k.pk, k.sk)),
           },
         },
       };
     } else {
       params = {
         RequestItems: {
-          [this.tableName]: {
+          [await this.tableName as string]: {
             Keys: keys.map((pk) => ({ [String(this.pk)]: pk })),
           },
         },
@@ -491,7 +529,7 @@ export default abstract class Model<T> {
       Object.assign(params, options);
     }
     const result = await this.documentClient.send(new BatchGetCommand(params));
-    return result.Responses ? (result.Responses[this.tableName] as T[]) : [];
+    return result.Responses ? (result.Responses[await this.tableName as string] as T[]) : [];
   }
 
   /**
@@ -565,7 +603,7 @@ export default abstract class Model<T> {
       updateActions['updatedAt'] = put(new Date().toISOString());
     }
     const params: UpdateCommandInput = {
-      TableName: this.tableName,
+      TableName: await this.tableName,
       Key: this.buildKeys(pk, sk),
       AttributeUpdates: buildUpdateActions(updateActions),
     };
@@ -671,10 +709,10 @@ export default abstract class Model<T> {
     //Make one BatchWrite request for every batch of 25 operations
     let params: BatchWriteCommandInput;
     const output: BatchWriteCommandOutput[] = await Promise.all(
-      batches.map(batch => {
+      batches.map(async (batch) => {
         params = {
           RequestItems: {
-            [this.tableName as string]: batch
+            [await this.tableName as string]: batch
           }
         }
         if (options) {

src/base-model.ts

@@ -20,6 +20,7 @@ export default class Query<T> extends Operation<T> {
   constructor(
     documentClient: DynamoDBDocumentClient,
     params: QueryCommandInput | ScanCommandInput,
+    private readonly tableName: Promise<string | undefined> | string | undefined,
     pk: string,
     sk?: string,
   ) {
@@ -92,6 +93,7 @@ export default class Query<T> extends Operation<T> {
   }
 
   public async doExec(): Promise<IPaginatedResult<T>> {
+    this.params.TableName = await this.tableName;
     const result = await this.documentClient.send(new QueryCommand(this.params));
     return this.buildResponse(result);
   }

src/query.ts

@@ -15,6 +15,7 @@ export default class Scan<T> extends Operation<T> {
   constructor(
     documentClient: DynamoDBDocumentClient,
     params: QueryCommandInput | ScanCommandInput,
+    private readonly tableName: Promise<string | undefined> | string | undefined,
     pk: string,
     sk?: string,
   ) {
@@ -47,6 +48,7 @@ export default class Scan<T> extends Operation<T> {
    * @returns Fetched items, and pagination metadata
    */
   public async doExec(): Promise<IPaginatedResult<T>> {
+    this.params.TableName = await this.tableName;
     const result = await this.documentClient.send(new ScanCommand(this.params));
     return this.buildResponse(result);
   }

src/scan.ts

@@ -0,0 +1,11 @@
+import { HashKeyEntity } from "./hashkey";
+import documentClient from './common';
+;
+import Model from "../../src/base-model";
+export default class SSMParamModel extends Model<HashKeyEntity> {
+    protected tableName = 'arn:aws:ssm:us-east-1:617599655210:parameter/tableName';
+
+    protected pk = 'hashkey';
+
+    protected documentClient = documentClient;
+}

test/models/ssm-param.ts

@@ -0,0 +1,75 @@
+import { SSMClient, GetParameterCommand } from '@aws-sdk/client-ssm';
+import { clearTables } from './hooks/create-tables';
+import * as modelApi from '../src/base-model';
+import SSMParamModel from './models/ssm-param';
+
+jest.mock('@aws-sdk/client-ssm');
+
+describe('fetchTableName', () => {
+    afterEach(async () => {
+        jest.clearAllMocks();
+    });
+    test('should throw an error', async () => {
+        const mockSend = jest.fn();
+        SSMClient.prototype.send = mockSend;
+        mockSend.mockRejectedValueOnce(new Error());
+        try {
+            await modelApi.fetchTableName('arn:aws:ssm:us-east-1:617599655210:parameter/tableName');
+        } catch (e) {
+            expect((e as Error).message.includes('Invalid SSM Parameter')).toBe(true);
+        }
+    });
+    test('should fetch tableName from SSM parameter store', async () => {
+        const mockSend = jest.fn();
+        mockSend.mockResolvedValueOnce({ Parameter: { Value: 'table_test_hashkey' } });
+        SSMClient.prototype.send = mockSend;
+        const tableName = await modelApi.fetchTableName('arn:aws:ssm:us-east-1:617599655210:parameter/tableName');
+        expect(GetParameterCommand).toHaveBeenCalledWith({ Name: 'tableName' });
+        expect(mockSend).toHaveBeenCalledWith(expect.any(GetParameterCommand));
+        expect(tableName).toEqual('table_test_hashkey');
+    });
+    test('should return the provided tableName', async () => {
+        const mockSend = jest.fn();
+        SSMClient.prototype.send = mockSend;
+        const mockResponse = { Parameter: { Value: 'table_test_hashkey' } };
+        mockSend.mockResolvedValueOnce(mockResponse);
+        const tableName = await modelApi.fetchTableName('tableName');
+        expect(mockSend).not.toHaveBeenCalled();
+        expect(GetParameterCommand).not.toHaveBeenCalled();
+        expect(tableName).toEqual('tableName');
+    });
+});
+
+describe('SSM parameter ARN', () => {
+    const item = {
+        hashkey: 'bar',
+        string: 'whatever',
+        stringmap: { foo: 'bar' },
+        stringset: ['bar, bar'],
+        number: 43,
+        bool: true,
+        list: ['foo', 42],
+    };
+    beforeEach(async () => {
+        const mockSend = jest.fn();
+        mockSend.mockResolvedValueOnce({ Parameter: { Value: 'table_test_hashkey' } });
+        SSMClient.prototype.send = mockSend;
+        await clearTables();
+    });
+    afterEach(async () => {
+        jest.clearAllMocks();
+    });
+    test('Should fetch tableName and save the item', async () => {
+        const model = new SSMParamModel();
+        await model.save(item);
+        const saved = await model.get('bar');
+        expect(saved).toEqual(item);
+    });
+    test('Should fetch tableName once and cache its value for subsequent requests', async () => {
+        jest.spyOn(modelApi, 'fetchTableName');
+        const model = new SSMParamModel();
+        await model.save(item);
+        await model.save(item);
+        expect(modelApi.fetchTableName).toHaveBeenCalledTimes(1);
+    });
+});