Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 25 additions & 8 deletions src/hope/contrib/aurora/celery_tasks.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from datetime import timedelta
import logging
from typing import TYPE_CHECKING, Any, Optional

Expand Down Expand Up @@ -125,14 +126,30 @@ def automate_rdi_creation_task(
@app.task
@log_start_and_end
@sentry_tags
def clean_old_record_files_task() -> None:
"""Task to clean (sets to null) Record's files field."""
from datetime import timedelta

def clean_old_record_files_task(batch_size: int = 100) -> None:
"""Task to remove old imported aurora records."""
try:
time_threshold = timezone.now() - timedelta(config.CLEARING_RECORD_FILES_TIMEDELTA)
Record.objects.filter(timestamp__lt=time_threshold, status=Record.STATUS_IMPORTED).delete()
logger.info("Record's files have benn successfully cleared")
except Exception as e:
logger.warning(e)
qs = Record.objects.filter(timestamp__lt=time_threshold, status=Record.STATUS_IMPORTED)
total = qs.count()
total_batches = (total + batch_size - 1) // batch_size

batch_num = 0
total_deleted = 0

while True:
batch = list(qs.values_list("pk", flat=True)[:batch_size])
if not batch:
break

deleted_count, _ = Record.objects.filter(pk__in=batch).delete()
batch_num += 1
total_deleted += deleted_count

logger.info(f"Batch {batch_num}/{total_batches} ({deleted_count} deleted, total: {total_deleted}/{total})")

logger.info("Record files have been successfully cleared")

except Exception: # noqa
logger.exception("Error cleaning old record files")
raise
46 changes: 46 additions & 0 deletions tests/unit/apps/aurora/test_clean_old_record_files_task.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from datetime import timedelta
from typing import TYPE_CHECKING

from django.utils import timezone
import pytest
Expand All @@ -7,6 +8,9 @@
from hope.contrib.aurora.celery_tasks import clean_old_record_files_task
from hope.contrib.aurora.models import Record

if TYPE_CHECKING:
from pytest_mock import MockerFixture

pytestmark = pytest.mark.django_db


Expand Down Expand Up @@ -50,3 +54,45 @@ def test_clean_old_record_files_task(record_set: dict[str, Record]) -> None:
assert record_set["recent_imported"].id in remaining_ids
assert record_set["old_error"].id in remaining_ids
assert record_set["old_to_import"].id in remaining_ids


def test_clean_old_record_files_task_empty() -> None:
# No records at all
clean_old_record_files_task()
assert Record.objects.count() == 0


def test_clean_old_record_files_task_batching(mocker: "MockerFixture") -> None:
now = timezone.now()
# Create 5 records that should be deleted
RecordFactory.create_batch(5, status=Record.STATUS_IMPORTED, timestamp=now - timedelta(days=100))

# Run with batch_size=2, should take 3 batches (2+2+1)
clean_old_record_files_task(batch_size=2)

assert Record.objects.count() == 0


def test_clean_old_record_files_task_logging(mocker: "MockerFixture") -> None:
now = timezone.now()
RecordFactory.create_batch(3, status=Record.STATUS_IMPORTED, timestamp=now - timedelta(days=100))
mock_logger = mocker.patch("hope.contrib.aurora.celery_tasks.logger")

clean_old_record_files_task(batch_size=2)

# Should log 2 batches and one final message
assert mock_logger.info.call_count == 3
mock_logger.info.assert_any_call("Batch 1/2 (2 deleted, total: 2/3)")
mock_logger.info.assert_any_call("Batch 2/2 (1 deleted, total: 3/3)")
mock_logger.info.assert_any_call("Record files have been successfully cleared")


def test_clean_old_record_files_task_error_handling(mocker: "MockerFixture") -> None:
mock_logger = mocker.patch("hope.contrib.aurora.celery_tasks.logger")
# Force an exception by mocking timezone.now to raise something
mocker.patch("hope.contrib.aurora.celery_tasks.timezone.now", side_effect=Exception("Database error"))

with pytest.raises(Exception, match="Database error"):
clean_old_record_files_task()

mock_logger.exception.assert_called_once_with("Error cleaning old record files")
Loading