Adding support for autoscaling based on a table tag as well as custom per-table config.

flow/aws-sdk.js

@@ -62,6 +62,16 @@ declare module 'aws-sdk' {
      TableNames: string[]
   };
 
+  declare type ListTagsRequest = {
+     NextToken?: string,
+     ResourceArn?: string
+  };
+
+  declare type ListTagsResponse = {
+     NextToken?: string,
+     Tags: string[]
+  };
+
   declare type DeleteTableRequest = {
      TableName: string,
   };

src/Provisioner.js

@@ -6,6 +6,7 @@ import Throughput from './utils/Throughput';
 import ProvisionerLogging from './provisioning/ProvisionerLogging';
 import { Region } from './configuration/Region';
 import DefaultProvisioner from './configuration/DefaultProvisioner';
+import CustomProvisioners from './configuration/CustomProvisioners';
 import { invariant } from './Global';
 import type { TableProvisionedAndConsumedThroughput, ProvisionerConfig, AdjustmentContext } from './flow/FlowTypes';
 
@@ -19,22 +20,54 @@ export default class Provisioner extends ProvisionerConfigurableBase {
   // Gets the list of tables which we want to autoscale
   async getTableNamesAsync(): Promise<string[]> {
 
-    // Option 1 - All tables (Default)
+    // Option 1: Identify tables by custom tag
+    if (process.env.DDB_AUTOSCALE_USE_TAGS) {
+      if (!process.env.AWS_REGION || !process.env.AWS_ACCOUNT_NUMBER || !process.env.DDB_AUTOSCALE_TAG_NAME) {
+        throw new Error('Missing environemnt variables to build the AWS ARN');
+      }
+
+      return await this.db.listAllTableNamesAsync()
+        .then(list => {
+          return Promise.all(list.map(name => {
+            const params = { ResourceArn: `arn:aws:dynamodb:${process.env.AWS_REGION}:${process.env.AWS_ACCOUNT_NUMBER}:table/${name}` };
+
+            return new Promise((resolve, reject) => {
+              // Required to throttle the requests (AWS only accepts 10 of these calls per second per account)
+              setTimeout(() => {
+                this.db.listTagsOfResourceAsync(params)
+                  .then(tags => resolve({ tableName: name, tags }))
+                  .catch(reject);
+              }, 100);
+            });
+          }));
+        })
+        .then(tableNamesWithTags => {
+          return tableNamesWithTags
+            .filter(pkg => { return pkg.tags.some(tag => tag.Key === process.env.DDB_AUTOSCALE_TAG_NAME || 'autoscaled' && tag.Value.match(/true/g)); })
+            .map(pkg => pkg.tableName);
+        })
+        .then(tableNames => {
+          ProvisionerLogging.logIdentifiedTables(tableNames);
+          return tableNames;
+        });
+    }
+
+    // Option 2 - All tables (Default)
     return await this.db.listAllTableNamesAsync();
 
-    // Option 2 - Hardcoded list of tables
+    // Option 3 - Hardcoded list of tables
     // return ['Table1', 'Table2', 'Table3'];
 
-    // Option 3 - DynamoDB / S3 configured list of tables
+    // Option 4 - DynamoDB / S3 configured list of tables
     // return await ...;
   }
 
   // Gets the json settings which control how the specifed table will be autoscaled
   // eslint-disable-next-line no-unused-vars
   getTableConfig(data: TableProvisionedAndConsumedThroughput): ProvisionerConfig {
 
-    // Option 1 - Default settings for all tables
-    return DefaultProvisioner;
+    // Option 1 - Default settings for all tables unless included in CustomProvisioners.json
+    return (CustomProvisioners || {})[data.TableName] || DefaultProvisioner;
 
     // Option 2 - Bespoke table specific settings
     // return data.TableName === 'Table1' ? Climbing : Default;

src/aws/DynamoDB.js

@@ -11,6 +11,8 @@ import type {
   UpdateTableResponse,
   ListTablesRequest,
   ListTablesResponse,
+  ListTagsRequest,
+  ListTagsResponse,
 } from 'aws-sdk';
 
 export default class DynamoDB {
@@ -52,6 +54,21 @@ export default class DynamoDB {
     }
   }
 
+  async listTagsAsync(params: ?ListTagsRequest): Promise<ListTagsResponse> {
+    let sw = stats.timer('DynamoDB.listTagsAsync').start();
+    try {
+      return await this._db.listTagsOfResource(params).promise();
+    } catch (ex) {
+      warning(JSON.stringify({
+        class: 'DynamoDB',
+        function: 'listTagsAsync'
+      }, null, json.padding));
+      throw ex;
+    } finally {
+      sw.end();
+    }
+  }
+
   async listAllTableNamesAsync(): Promise<string[]> {
     let tableNames = [];
     let lastTable;
@@ -63,6 +80,17 @@ export default class DynamoDB {
     return tableNames;
   }
 
+  async listTagsOfResourceAsync(params): Promise<string[]> {
+    let tagValues = [];
+    let nextToken;
+    do {
+      let listTagsResponse = await this.listTagsAsync(params);
+      tagValues = tagValues.concat(listTagsResponse.Tags);
+      nextToken = listTagsResponse.NextToken;
+    } while (nextToken);
+    return tagValues;
+  }
+
   async describeTableAsync(params: DescribeTableRequest): Promise<DescribeTableResponse> {
     let sw = stats.timer('DynamoDB.describeTableAsync').start();
     try {

src/provisioning/ProvisionerLogging.js

@@ -6,6 +6,10 @@ import type {
 } from '../flow/FlowTypes';
 
 export default class ConfigLogging {
+  static logIdentifiedTables(tableNames) {
+    log('The following tables were identified for autoscaling:', tableNames);
+  }
+
   static isAdjustmentRequiredLog(
     adjustmentContext: AdjustmentContext,
     adjustmentData: AdjustmentData,