fix purge table when table has many keys

aaronshaf · Mar 1, 2019 · d7f27b0 · d7f27b0
1 parent 8660dd9
commit d7f27b0
Showing 1 changed file with 13 additions and 9 deletions.
diff --git 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 }