Fix/restore dynamodump from backup with gsis (#94)

* Checks provided read and write capacity units in schema to be greater than 0 for table and GSIs, otherwise fallback to default RESTORE_READ_CAPACITY to avoid the process crashing

* Adds functions to prepare payload passed to create_table boto3 API call in on_restore function

* Adds .idea to gitignore

* Fixes typo in original_gsi_read_capacity

* Lints files

* Fix failing tests

* Lints files

Author: w0rks74710n

.gitignore

@@ -39,3 +39,4 @@ dump
 dynamodump.iml
 env
 .vscode
+.idea

dynamodump/dynamodump.py

@@ -44,6 +44,7 @@
 MAX_RETRY = 6
 METADATA_URL = "http://169.254.169.254/latest/meta-data/"
 RESTORE_WRITE_CAPACITY = 25
+RESTORE_READ_CAPACITY = 25
 SCHEMA_FILE = "schema.json"
 THREAD_START_DELAY = 1  # seconds
 
@@ -747,6 +748,33 @@ def do_backup(dynamo, read_capacity, tableQueue=None, srcTable=None):
             tableQueue.task_done()
 
 
+def prepare_provisioned_throughput_for_restore(provisioned_throughput):
+    """
+    This function makes sure that the payload returned for the boto3 API call create_table is compatible
+    with the provisioned throughput attribute
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html
+    """
+    return {
+        "ReadCapacityUnits": provisioned_throughput["ReadCapacityUnits"],
+        "WriteCapacityUnits": provisioned_throughput["WriteCapacityUnits"],
+    }
+
+
+def prepare_gsi_for_restore(gsi):
+    """
+    This function makes sure that the payload returned for the boto3 API call create_table is compatible
+    See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb.html
+    """
+    return {
+        "IndexName": gsi["IndexName"],
+        "KeySchema": gsi["KeySchema"],
+        "Projection": gsi["Projection"],
+        "ProvisionedThroughput": prepare_provisioned_throughput_for_restore(
+            gsi["ProvisionedThroughput"]
+        ),
+    }
+
+
 def do_restore(dynamo, sleep_interval, source_table, destination_table, write_capacity):
     """
     Restore table
@@ -792,21 +820,58 @@ def do_restore(dynamo, sleep_interval, source_table, destination_table, write_ca
         else:
             write_capacity = original_write_capacity
 
+    if original_write_capacity == 0:
+        original_write_capacity = RESTORE_WRITE_CAPACITY
+
+    # ensure that read capacity is at least RESTORE_READ_CAPACITY
+    if original_read_capacity < RESTORE_READ_CAPACITY:
+        read_capacity = RESTORE_WRITE_CAPACITY
+    else:
+        read_capacity = original_read_capacity
+
+    if original_read_capacity == 0:
+        original_read_capacity = RESTORE_READ_CAPACITY
+
     # override GSI write capacities if specified, else use RESTORE_WRITE_CAPACITY if original
     # write capacity is lower
     original_gsi_write_capacities = []
+    original_gsi_read_capacities = []
     if table_global_secondary_indexes is not None:
         for gsi in table_global_secondary_indexes:
-            original_gsi_write_capacities.append(
-                gsi["ProvisionedThroughput"]["WriteCapacityUnits"]
-            )
+            # keeps track of original gsi write capacity units. If provisioned capacity is 0, set to
+            # RESTORE_WRITE_CAPACITY as fallback given that 0 is not allowed for write capacities
+            original_gsi_write_capacity = gsi["ProvisionedThroughput"][
+                "WriteCapacityUnits"
+            ]
+            if original_gsi_write_capacity == 0:
+                original_gsi_write_capacity = RESTORE_WRITE_CAPACITY
+
+            original_gsi_write_capacities.append(original_gsi_write_capacity)
 
             if gsi["ProvisionedThroughput"]["WriteCapacityUnits"] < int(write_capacity):
                 gsi["ProvisionedThroughput"]["WriteCapacityUnits"] = int(write_capacity)
 
+            # keeps track of original gsi read capacity units. If provisioned capacity is 0, set to
+            # RESTORE_READ_CAPACITY as fallback given that 0 is not allowed for read capacities
+            original_gsi_read_capacity = gsi["ProvisionedThroughput"][
+                "ReadCapacityUnits"
+            ]
+            if original_gsi_read_capacity == 0:
+                original_gsi_read_capacity = RESTORE_READ_CAPACITY
+
+            original_gsi_read_capacities.append(original_gsi_read_capacity)
+
+            if (
+                gsi["ProvisionedThroughput"]["ReadCapacityUnits"]
+                < RESTORE_READ_CAPACITY
+            ):
+                gsi["ProvisionedThroughput"][
+                    "ReadCapacityUnits"
+                ] = RESTORE_READ_CAPACITY
+
     # temp provisioned throughput for restore
     table_provisioned_throughput = {
-        "ReadCapacityUnits": int(original_read_capacity),
+        "ReadCapacityUnits": int(read_capacity),
         "WriteCapacityUnits": int(write_capacity),
     }
 
@@ -824,7 +889,9 @@ def do_restore(dynamo, sleep_interval, source_table, destination_table, write_ca
             optional_args["LocalSecondaryIndexes"] = table_local_secondary_indexes
 
         if table_global_secondary_indexes is not None:
-            optional_args["GlobalSecondaryIndexes"] = table_global_secondary_indexes
+            optional_args["GlobalSecondaryIndexes"] = [
+                prepare_gsi_for_restore(gsi) for gsi in table_global_secondary_indexes
+            ]
 
         while True:
             try:
@@ -914,7 +981,10 @@ def do_restore(dynamo, sleep_interval, source_table, destination_table, write_ca
 
         if not args.skipThroughputUpdate:
             # revert to original table write capacity if it has been modified
-            if int(write_capacity) != original_write_capacity:
+            if (
+                int(write_capacity) != original_write_capacity
+                or int(read_capacity) != original_read_capacity
+            ):
                 update_provisioned_throughput(
                     dynamo,
                     destination_table,
@@ -930,13 +1000,19 @@ def do_restore(dynamo, sleep_interval, source_table, destination_table, write_ca
                     wcu = gsi["ProvisionedThroughput"]["WriteCapacityUnits"]
                     rcu = gsi["ProvisionedThroughput"]["ReadCapacityUnits"]
                     original_gsi_write_capacity = original_gsi_write_capacities.pop(0)
-                    if original_gsi_write_capacity != wcu:
+                    original_gsi_read_capacity = original_gsi_read_capacities.pop(0)
+                    if (
+                        original_gsi_write_capacity != wcu
+                        or original_gsi_read_capacity != rcu
+                    ):
                         gsi_data.append(
                             {
                                 "Update": {
                                     "IndexName": gsi["IndexName"],
                                     "ProvisionedThroughput": {
-                                        "ReadCapacityUnits": int(rcu),
+                                        "ReadCapacityUnits": int(
+                                            original_gsi_read_capacity
+                                        ),
                                         "WriteCapacityUnits": int(
                                             original_gsi_write_capacity
                                         ),
@@ -945,30 +1021,31 @@ def do_restore(dynamo, sleep_interval, source_table, destination_table, write_ca
                             }
                         )
 
-                logging.info(
-                    "Updating "
-                    + destination_table
-                    + " global secondary indexes write capacities as necessary.."
-                )
-                while True:
-                    try:
-                        dynamo.update_table(
-                            TableName=destination_table,
-                            GlobalSecondaryIndexUpdates=gsi_data,
-                        )
-                        break
-                    except dynamo.exceptions.LimitExceededException:
-                        logging.info(
-                            "Limit exceeded, retrying updating throughput of"
-                            "GlobalSecondaryIndexes in " + destination_table + ".."
-                        )
-                        time.sleep(sleep_interval)
-                    except dynamo.exceptions.ProvisionedThroughputExceededException:
-                        logging.info(
-                            "Control plane limit exceeded, retrying updating throughput of"
-                            "GlobalSecondaryIndexes in " + destination_table + ".."
-                        )
-                        time.sleep(sleep_interval)
+                if gsi_data:
+                    logging.info(
+                        "Updating "
+                        + destination_table
+                        + " global secondary indexes write and read capacities as necessary.."
+                    )
+                    while True:
+                        try:
+                            dynamo.update_table(
+                                TableName=destination_table,
+                                GlobalSecondaryIndexUpdates=gsi_data,
+                            )
+                            break
+                        except dynamo.exceptions.LimitExceededException:
+                            logging.info(
+                                "Limit exceeded, retrying updating throughput of"
+                                "GlobalSecondaryIndexes in " + destination_table + ".."
+                            )
+                            time.sleep(sleep_interval)
+                        except dynamo.exceptions.ProvisionedThroughputExceededException:
+                            logging.info(
+                                "Control plane limit exceeded, retrying updating throughput of"
+                                "GlobalSecondaryIndexes in " + destination_table + ".."
+                            )
+                            time.sleep(sleep_interval)
 
         # wait for table to become active
         wait_for_active_table(dynamo, destination_table, "active")