diff --git a/flow/aws-sdk.js b/flow/aws-sdk.js index 4ebae82..c3e391f 100644 --- a/flow/aws-sdk.js +++ b/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, }; diff --git a/src/Provisioner.js b/src/Provisioner.js index e214b0d..610b3bf 100644 --- a/src/Provisioner.js +++ b/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,13 +20,45 @@ export default class Provisioner extends ProvisionerConfigurableBase { // Gets the list of tables which we want to autoscale async getTableNamesAsync(): Promise { - // 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 ...; } @@ -33,8 +66,8 @@ export default class Provisioner extends ProvisionerConfigurableBase { // 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; diff --git a/src/aws/DynamoDB.js b/src/aws/DynamoDB.js index e46174f..ce20549 100644 --- a/src/aws/DynamoDB.js +++ b/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 { + 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 { let tableNames = []; let lastTable; @@ -63,6 +80,17 @@ export default class DynamoDB { return tableNames; } + async listTagsOfResourceAsync(params): Promise { + 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 { let sw = stats.timer('DynamoDB.describeTableAsync').start(); try { diff --git a/src/provisioning/ProvisionerLogging.js b/src/provisioning/ProvisionerLogging.js index 2bb4bc8..872abec 100644 --- a/src/provisioning/ProvisionerLogging.js +++ b/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,