From 3987f6099b85affc094f4b2ead911c6ca7d13737 Mon Sep 17 00:00:00 2001 From: Rowan Seymour Date: Tue, 20 Aug 2024 23:09:32 +0000 Subject: [PATCH] Tweak migrate_dynamo command --- .../management/commands/migrate_dynamo.py | 28 +++++++++++-------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/temba/utils/management/commands/migrate_dynamo.py b/temba/utils/management/commands/migrate_dynamo.py index 9679103142e..724112d0ad5 100644 --- a/temba/utils/management/commands/migrate_dynamo.py +++ b/temba/utils/management/commands/migrate_dynamo.py @@ -18,7 +18,7 @@ class Command(BaseCommand): help = "Creates DynamoDB tables that don't already exist." def handle(self, *args, **kwargs): - client = dynamo.get_client() + self.client = dynamo.get_client() for table in TABLES: # add optional prefix to name to allow multiple deploys in same region @@ -28,17 +28,21 @@ def handle(self, *args, **kwargs): # ttl isn't actually part of the create_table call ttlSpec = table.pop("TimeToLiveSpecification", None) - try: - client.describe_table(TableName=table["TableName"]) + if not self._table_exists(name): + self.client.create_table(**table) - self.stdout.write(f"{table['TableName']}: already exists") - except client.exceptions.ResourceNotFoundException: - client.create_table(**table) + self.stdout.write(f"{name}: created") if ttlSpec: - client.update_time_to_live( - TableName=table["TableName"], - TimeToLiveSpecification=ttlSpec, - ) - - self.stdout.write(f"{table['TableName']}: created") + self.client.update_time_to_live(TableName=name, TimeToLiveSpecification=ttlSpec) + + self.stdout.write(f"{name}: updated TTL") + else: + self.stdout.write(f"{name}: already exists") + + def _table_exists(self, name: str) -> bool: + try: + self.client.describe_table(TableName=name) + return True + except self.client.exceptions.ResourceNotFoundException: + return False