From 98d8e4ec93e1715c7ed31a6a206beb2f11ee955e Mon Sep 17 00:00:00 2001 From: David Kennedy Date: Tue, 30 Jul 2024 22:10:05 -0400 Subject: [PATCH] updates --- src/registrar/management/commands/clean_tables.py | 11 +++++++---- src/registrar/tests/test_management_scripts.py | 4 ---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/registrar/management/commands/clean_tables.py b/src/registrar/management/commands/clean_tables.py index dd84b0b62..5d4439d95 100644 --- a/src/registrar/management/commands/clean_tables.py +++ b/src/registrar/management/commands/clean_tables.py @@ -64,15 +64,18 @@ class Command(BaseCommand): model = apps.get_model("registrar", table_name) BATCH_SIZE = 1000 total_deleted = 0 - while True: - pks = list(model.objects.values_list("pk", flat=True)[:BATCH_SIZE]) - if not pks: - break + + # Get initial batch of primary keys + pks = list(model.objects.values_list("pk", flat=True)[:BATCH_SIZE]) + + while pks: # Use a transaction to ensure database integrity with transaction.atomic(): deleted, _ = model.objects.filter(pk__in=pks).delete() total_deleted += deleted logger.debug(f"Deleted {deleted} {table_name}s, total deleted: {total_deleted}") + # Get the next batch of primary keys + pks = list(model.objects.values_list("pk", flat=True)[:BATCH_SIZE]) logger.info(f"Successfully cleaned table {table_name}, deleted {total_deleted} rows") except LookupError: logger.error(f"Model for table {table_name} not found.") diff --git a/src/registrar/tests/test_management_scripts.py b/src/registrar/tests/test_management_scripts.py index 8fb754ba0..bfdeefadd 100644 --- a/src/registrar/tests/test_management_scripts.py +++ b/src/registrar/tests/test_management_scripts.py @@ -33,10 +33,6 @@ from api.tests.common import less_console_noise_decorator logger = logging.getLogger(__name__) -class CustomDeleteException(Exception): - pass - - class TestPopulateVerificationType(MockEppLib): """Tests for the populate_organization_type script"""