APM-3095 Paginate on force unlock

Author: acarriedev

scripts/terraform_force_unlock.py

@@ -27,40 +27,69 @@
 
 @click.command()
 @click.option("--min-age-hr", type=int, default=4)
-@click.option("--key-prefix", type=str, default="nhsd-apm-management-ptl-terraform/env:/api-deployment:ptl:")
+@click.option(
+    "--key-prefix",
+    type=str,
+    default="nhsd-apm-management-ptl-terraform/env:/api-deployment:ptl:",
+)
 @click.option("--table-name", type=str, default="terraform-state-lock")
 @click.option("--profile", type=str, default="apm_ptl")
 def main(min_age_hr, key_prefix, table_name, profile):
-
     accepted_envs = ["apm_ptl", "apm_prod"]
 
     if profile not in accepted_envs:
         raise ValueError("Profile must be apm_ptl or apm_prod")
 
-    terraform_lock_table = boto3.Session(profile_name=profile).resource("dynamodb").Table(table_name)
+    terraform_lock_table = (
+        boto3.Session(profile_name=profile).resource("dynamodb").Table(table_name)
+    )
 
     filter_expr = "begins_with(#n0, :v0) AND attribute_exists(#n1)"
 
     ExpressionAttributeNames = {"#n0": "LockID", "#n1": "Info"}
     ExpressionAttributeValues = {
         ":v0": key_prefix,
     }
-    items = terraform_lock_table.scan(FilterExpression=filter_expr, ExpressionAttributeNames=ExpressionAttributeNames, ExpressionAttributeValues=ExpressionAttributeValues)
-    print(f"Found {len(items['Items'])} locks which start with key prefix '{key_prefix}'")
+    items = terraform_lock_table.scan(
+        FilterExpression=filter_expr,
+        ExpressionAttributeNames=ExpressionAttributeNames,
+        ExpressionAttributeValues=ExpressionAttributeValues,
+    )
+
+    total_items = items["Items"]
+
+    while "LastEvaluatedKey" in items:
+        items = terraform_lock_table.scan(
+            FilterExpression=filter_expr,
+            ExpressionAttributeNames=ExpressionAttributeNames,
+            ExpressionAttributeValues=ExpressionAttributeValues,
+            ExclusiveStartKey=items["LastEvaluatedKey"],
+        )
+        total_items.extend(items["Items"])
+
+    print(
+        f"Found {len(items['Items'])} locks which start with key prefix '{key_prefix}'"
+    )
 
     removed_count = 0
-    for lock_item in items["Items"]:
+    for lock_item in total_items:
         lock_item_info = json.loads(lock_item["Info"])
         lock_id = lock_item["LockID"]
         created_at = dateutil.parser.parse(lock_item_info["Created"])
 
-        if datetime.datetime.now(datetime.timezone.utc) - created_at > datetime.timedelta(hours=min_age_hr):
-            print(f"{lock_id} {created_at=} is more than {min_age_hr} hours old, deleting lock...")
+        if datetime.datetime.now(
+            datetime.timezone.utc
+        ) - created_at > datetime.timedelta(hours=min_age_hr):
+            print(
+                f"{lock_id} {created_at=} is more than {min_age_hr} hours old, deleting lock..."
+            )
             terraform_lock_table.delete_item(Key={"LockID": lock_id})
             removed_count += 1
 
         else:
-            print(f"{lock_id} {created_at=} is not more than {min_age_hr} hours old, leaving it alone!")
+            print(
+                f"{lock_id} {created_at=} is not more than {min_age_hr} hours old, leaving it alone!"
+            )
 
     print(f"Removed {removed_count} locks")