From ec873319530abceaf1377325449f0d1a8a10de29 Mon Sep 17 00:00:00 2001
From: Elad Ben-Israel <benisrae@amazon.com>
Date: Wed, 31 Oct 2018 16:45:14 +0200
Subject: [PATCH] feat(aws-dynamodb): allow specifying partition/sort keys in
 props (#1054)

Adds an option to specify `partitionKey` and/or `sortKey` when the
table is initialized. This is identical to calling `addPartitionKey`
and `addSortKey`.

Fixes #1051
---
 packages/@aws-cdk/aws-dynamodb/README.md      | 16 ++++-----
 packages/@aws-cdk/aws-dynamodb/lib/table.ts   | 23 +++++++++++--
 .../aws-dynamodb/test/test.dynamodb.ts        | 34 ++++++++++++++++++-
 3 files changed, 62 insertions(+), 11 deletions(-)

diff --git a/packages/@aws-cdk/aws-dynamodb/README.md b/packages/@aws-cdk/aws-dynamodb/README.md
index 73b2b2dc91e0b..b5174f9501690 100644
--- a/packages/@aws-cdk/aws-dynamodb/README.md
+++ b/packages/@aws-cdk/aws-dynamodb/README.md
@@ -1,20 +1,20 @@
 ## AWS DynamoDB Construct Library
 
-Add a DynamoDB table to your stack like so:
+Here is a minimal deployable DynamoDB table definition:
 
 ```ts
 import dynamodb = require('@aws-cdk/aws-dynamodb');
 
 const table = new dynamodb.Table(stack, 'Table', {
-    // You can leave this out to automatically generate a name.
-    tableName: 'MyTableName',
-
-    // If you leave these out they default to 5
-    readCapacity: 100,
-    writeCapacity: 10,
-})
+  partitionKey: { name: 'id', type: dynamodb.AttributeType.String }
+});
 ```
 
+### Keys
+
+You can either specify `partitionKey` and/or `sortKey` when you initialize the
+table, or call `addPartitionKey` and `addSortKey` after initialization.
+
 ### Configure AutoScaling for your table
 
 You can have DynamoDB automatically raise and lower the read and write capacities
diff --git a/packages/@aws-cdk/aws-dynamodb/lib/table.ts b/packages/@aws-cdk/aws-dynamodb/lib/table.ts
index 948e078fdbf5f..783718303cea0 100644
--- a/packages/@aws-cdk/aws-dynamodb/lib/table.ts
+++ b/packages/@aws-cdk/aws-dynamodb/lib/table.ts
@@ -87,6 +87,18 @@ export interface TableProps {
    * @default undefined, TTL is disabled
    */
   ttlAttributeName?: string;
+
+  /**
+   * Partition key attribute definition. This is eventually required, but you
+   * can also use `addPartitionKey` to specify the partition key at a later stage.
+   */
+  partitionKey?: Attribute;
+
+  /**
+   * Table sort key attribute definition. You can also use `addSortKey` to set
+   * this up later.
+   */
+  sortKey?: Attribute;
 }
 
 export interface SecondaryIndexProps {
@@ -158,8 +170,8 @@ export class Table extends Construct {
   private readonly secondaryIndexNames: string[] = [];
   private readonly nonKeyAttributes: string[] = [];
 
-  private tablePartitionKey: Attribute | undefined = undefined;
-  private tableSortKey: Attribute | undefined = undefined;
+  private tablePartitionKey?: Attribute;
+  private tableSortKey?: Attribute;
 
   private readonly tableScaling: ScalableAttributePair = {};
   private readonly indexScaling = new Map<string, ScalableAttributePair>();
@@ -190,6 +202,13 @@ export class Table extends Construct {
 
     this.scalingRole = this.makeScalingRole();
 
+    if (props.partitionKey) {
+      this.addPartitionKey(props.partitionKey);
+    }
+
+    if (props.sortKey) {
+      this.addSortKey(props.sortKey);
+    }
   }
 
   /**
diff --git a/packages/@aws-cdk/aws-dynamodb/test/test.dynamodb.ts b/packages/@aws-cdk/aws-dynamodb/test/test.dynamodb.ts
index 0f3b945c53389..e3cbc780c4d34 100644
--- a/packages/@aws-cdk/aws-dynamodb/test/test.dynamodb.ts
+++ b/packages/@aws-cdk/aws-dynamodb/test/test.dynamodb.ts
@@ -123,6 +123,38 @@ export = {
       test.done();
     },
 
+    'hash + range key can also be specified in props'(test: Test) {
+      const app = new TestApp();
+
+      new Table(app.stack, CONSTRUCT_NAME, {
+        partitionKey: TABLE_PARTITION_KEY,
+        sortKey: TABLE_SORT_KEY
+      });
+
+      const template = app.synthesizeTemplate();
+
+      test.deepEqual(template, {
+        Resources: {
+          MyTable794EDED1: {
+            Type: 'AWS::DynamoDB::Table',
+            Properties: {
+              AttributeDefinitions: [
+                { AttributeName: 'hashKey', AttributeType: 'S' },
+                { AttributeName: 'sortKey', AttributeType: 'N' }
+              ],
+              KeySchema: [
+                { AttributeName: 'hashKey', KeyType: 'HASH' },
+                { AttributeName: 'sortKey', KeyType: 'RANGE' }
+              ],
+              ProvisionedThroughput: { ReadCapacityUnits: 5, WriteCapacityUnits: 5 },
+            }
+          }
+        }
+      });
+
+      test.done();
+    },
+
     'point-in-time recovery is not enabled'(test: Test) {
       const app = new TestApp();
       new Table(app.stack, CONSTRUCT_NAME)
@@ -1167,7 +1199,7 @@ export = {
     '"grantFullAccess" allows the principal to perform any action on the table ("*")'(test: Test) {
       testGrant(test, [ '*' ], (p, t) => t.grantFullAccess(p));
     }
-  }
+  },
 };
 
 class TestApp {