From d7f27b09ed0e9ae7c217cca66e06ad980d82e90f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juan=20Andr=C3=A9s=20Carmena?= <juan14nob@gmail.com>
Date: Tue, 26 Feb 2019 10:54:08 -0300
Subject: [PATCH] fix purge table when table has many keys

---
 lib/actions/purgeTable.js | 22 +++++++++++++---------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/lib/actions/purgeTable.js b/lib/actions/purgeTable.js
index c5f5aaa..4b2ac6b 100644
--- a/lib/actions/purgeTable.js
+++ b/lib/actions/purgeTable.js
@@ -1,14 +1,20 @@
 const { doSearch } = require('../util')
 
-const findPrimaryKey = context => {
+const findPrimaryKeys = context => {
   const params = {
     TableName: context.tableName
   }
 
   return context.dynamodb.describeTable(params).promise()
     .then(tableDescription => {
-      context.primaryKey = tableDescription.Table.KeySchema.find(
-        element => element.KeyType === 'HASH').AttributeName
+      const tableSchema = tableDescription.Table.KeySchema
+
+      context.primaryKeys = ['HASH', 'RANGE']
+        .map(keyType => tableSchema.find(element => element.KeyType === keyType))
+        .filter(attribute => attribute)
+        .map(attribute => attribute.AttributeName)
+
+
       return context
     })
 }
@@ -26,9 +32,7 @@ const deleteAllElements = context => {
   for (const item of context.items) {
     params.RequestItems[context.tableName].push({
       DeleteRequest: {
-        Key: {
-          [context.primaryKey]: item[context.primaryKey]
-        }
+        Key: item
       }
     })
 
@@ -50,7 +54,7 @@ const deleteAllElements = context => {
 
 const findAllElements = context => {
   const scanParams = {
-    ProjectionExpression: context.primaryKey,
+    ProjectionExpression: context.primaryKeys.join(', '),
   }
 
   return doSearch(context.dynamodb, context.tableName, scanParams)
@@ -81,9 +85,9 @@ const purgeTable = (tableName, dynamodb) => {
     dynamodb
   }
 
-  return findPrimaryKey(context)
+  return findPrimaryKeys(context)
     .then(findAllElements)
     .then(deleteAllElements)
 }
 
-module.exports = {purgeTable}
+module.exports = { purgeTable }