Merge pull request #29 from maykinmedia/flush-cmd

Add a flush command

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",

tests/test_management_commands.py

@@ -1,4 +1,5 @@
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone as dt_timezone
+from io import StringIO
 
 from django.conf import settings
 from django.core import mail
@@ -7,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, UserFactory
+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]")
@@ -152,3 +156,46 @@ def test_timeline_digest_from_email_setting(self):
 
         self.assertEqual(len(mail.outbox), 1)
         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, 0, 0, 0, tzinfo=dt_timezone.utc)
+        self.log_1.save()
+
+        self.log_2 = TimelineLogFactory.create()
+        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",
+            "--all",
+            interactive=False,
+            verbosity=0,
+            stdout=stdout,
+        )
+
+        self.assertEqual(TimelineLog.objects.count(), 0)
+        self.assertEqual(
+            stdout.getvalue().strip(), "Successfully deleted 2 timeline logs."
+        )
+
+    def test_prune_timeline_logs_date(self):
+        call_command(
+            "prune_timeline_logs",
+            "--keep-days",
+            "2",
+            interactive=False,
+            verbosity=0,
+            stdout=StringIO(),
+        )
+
+        self.assertEqual(TimelineLog.objects.count(), 1)
+        self.assertEqual(TimelineLog.objects.first().pk, self.log_2.pk)

timeline_logger/management/commands/prune_timeline_logs.py

@@ -0,0 +1,54 @@
+from textwrap import dedent
+
+from django.core.management.base import BaseCommand
+
+from timeline_logger.service import prune_timeline_logs
+
+
+class Command(BaseCommand):
+    help = "Removes timeline logs objects older than the specified date."
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--noinput",
+            "--no-input",
+            action="store_false",
+            dest="interactive",
+            help="Tells Django to NOT prompt the user for input of any kind.",
+        )
+        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"]
+
+        if all and interactive:
+            confirm = input(
+                dedent(
+                    """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: """
+                )
+            )
+        else:
+            confirm = "yes"
+
+        if confirm == "yes":
+            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