Feedback

pyproject.toml

@@ -42,6 +42,7 @@ Changelog = "https://github.com/maykinmedia/django-timeline-logger/blob/master/d
 [project.optional-dependencies]
 tests = [
     "factory-boy",
+    "time-machine",
     "psycopg2",
     "pytest",
     "pytest-cov",
@@ -80,7 +81,7 @@ sections=["FUTURE", "STDLIB", "DJANGO", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
-DJANGO_SETTINGS_MODULE = "tests.settings_pg"
+DJANGO_SETTINGS_MODULE = "tests.settings"
 
 [tool.bumpversion]
 current_version = "4.0.0"

tests/test_management_commands.py

@@ -8,13 +8,16 @@
 from django.test import TestCase, override_settings
 from django.utils import timezone
 
+import time_machine
+
 from timeline_logger.models import TimelineLog
 
 from .factories import ArticleFactory, TimelineLogFactory, UserFactory
 
 
 class ReportMailingTestCase(TestCase):
     def setUp(self):
+        super().setUp()
         self.article = ArticleFactory.create()
 
         self.user = UserFactory.create(email="[email protected]")
@@ -155,33 +158,40 @@ def test_timeline_digest_from_email_setting(self):
         self.assertEqual(mail.outbox[0].from_email, settings.TIMELINE_DIGEST_FROM_EMAIL)
 
 
+@time_machine.travel(datetime(2024, 3, 5, 0, 0, 0, tzinfo=dt_timezone.utc))
 class PruneTimelineLogsTestCase(TestCase):
     def setUp(self):
+        super().setUp()
 
         self.log_1 = TimelineLogFactory.create()
-        self.log_1.timestamp = datetime(2024, 3, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
+        self.log_1.timestamp = datetime(2024, 3, 1, 0, 0, 0, tzinfo=dt_timezone.utc)
         self.log_1.save()
 
         self.log_2 = TimelineLogFactory.create()
-        self.log_2.timestamp = datetime(2024, 3, 1, 14, 0, 0, tzinfo=dt_timezone.utc)
+        self.log_2.timestamp = datetime(2024, 3, 4, 0, 0, 0, tzinfo=dt_timezone.utc)
         self.log_2.save()
 
     def test_prune_timeline_logs_no_date(self):
         stdout = StringIO()
 
         call_command(
-            "prune_timeline_logs", interactive=False, verbosity=0, stdout=stdout
+            "prune_timeline_logs",
+            "--all",
+            interactive=False,
+            verbosity=0,
+            stdout=stdout,
         )
 
         self.assertEqual(TimelineLog.objects.count(), 0)
-        stdout.seek(0)
-        self.assertEqual(stdout.read().strip(), "Successfully deleted 2 timeline logs.")
+        self.assertEqual(
+            stdout.getvalue().strip(), "Successfully deleted 2 timeline logs."
+        )
 
     def test_prune_timeline_logs_date(self):
         call_command(
             "prune_timeline_logs",
-            "--before",
-            "2024-03-01T13:00:00+00:00",
+            "--keep-days",
+            "2",
             interactive=False,
             verbosity=0,
             stdout=StringIO(),

timeline_logger/management/commands/prune_timeline_logs.py

@@ -1,9 +1,8 @@
-from datetime import datetime
 from textwrap import dedent
 
 from django.core.management.base import BaseCommand
 
-from timeline_logger.models import TimelineLog
+from timeline_logger.service import prune_timeline_logs
 
 
 class Command(BaseCommand):
@@ -17,21 +16,30 @@ def add_arguments(self, parser):
             dest="interactive",
             help="Tells Django to NOT prompt the user for input of any kind.",
         )
-        parser.add_argument(
-            "--before",
-            type=datetime.fromisoformat,
-            help="Only flush timeline logs older than the specified date.",
+        exclusive_group = parser.add_mutually_exclusive_group(required=True)
+
+        exclusive_group.add_argument(
+            "--all",
+            action="store_true",
+            help="Whether to delete all log records.",
+        )
+
+        exclusive_group.add_argument(
+            "--keep-days",
+            type=int,
+            help="Only delete records older than the specified number of days.",
         )
 
     def handle(self, *args, **options):
+        all = options["all"]
+        keep_days = options["keep_days"]
         interactive = options["interactive"]
-        before = options["before"]
 
-        if not before and interactive:
+        if all and interactive:
             confirm = input(
                 dedent(
-                    """You haven't specified a date to limit the objects to be deleted.
-                This will delete all timeline logs objects. Are you sure you want to do this?
+                    """You have specified "--all", meaning all timeline logs will be deleted.
+                Are you sure you want to do this?
 
                 Type 'yes' to continue, or 'no' to cancel: """
                 )
@@ -40,11 +48,7 @@ def handle(self, *args, **options):
             confirm = "yes"
 
         if confirm == "yes":
-            if before:
-                qs = TimelineLog.objects.filter(timestamp__lte=before)
-            else:
-                qs = TimelineLog.objects.all()
-            number, _ = qs.delete()
+            number = prune_timeline_logs(keep_days=0 if all else keep_days)
             self.stdout.write(f"Successfully deleted {number} timeline logs.")
         else:
             self.stdout.write("Flush cancelled.")

timeline_logger/service.py

@@ -0,0 +1,19 @@
+from datetime import timedelta
+
+from django.utils import timezone
+
+from .models import TimelineLog
+
+
+def prune_timeline_logs(*, keep_days: int | None = None) -> int:
+    """Delete the timeline logs instances.
+
+    :param keep_days: If specified, only delete records older than the specified number of days.
+    :returns: The number of deleted instances.
+    """
+    limit = timezone.now()
+    if keep_days is not None:
+        limit -= timedelta(days=keep_days)
+
+    number, _ = TimelineLog.objects.filter(timestamp__lte=limit).delete()
+    return number