feat: Add truncate option to flush command (#681)

* [feat]: Add truncate option to flush command

* [test]: Add test cases for flush command with truncate

* [refactor]: Simplified truncate query class and remove redundant return statement

* [test]: Add test cases to test truncate for unsupported database vendor

* [docs]: Update change log

Author: aabharya

CHANGELOG.md

@@ -5,6 +5,7 @@
 #### Improvements
 
 - feat: Added `LogEntry.remote_port` field. ([#671](https://github.com/jazzband/django-auditlog/pull/671))
+- feat: Added `truncate` option to `auditlogflush` management command. ([#681](https://github.com/jazzband/django-auditlog/pull/681))
 - Drop Python 3.8 support. ([#678](https://github.com/jazzband/django-auditlog/pull/678))
 - Confirm Django 5.1 support and drop Django 3.2 support. ([#677](https://github.com/jazzband/django-auditlog/pull/677))
 

auditlog/management/commands/auditlogflush.py

@@ -1,6 +1,7 @@
 import datetime
 
 from django.core.management.base import BaseCommand
+from django.db import connection
 
 from auditlog.models import LogEntry
 
@@ -25,11 +26,24 @@ def add_arguments(self, parser):
             dest="before_date",
             type=datetime.date.fromisoformat,
         )
+        parser.add_argument(
+            "-t",
+            "--truncate",
+            action="store_true",
+            default=None,
+            help="Truncate log entry table.",
+            dest="truncate",
+        )
 
     def handle(self, *args, **options):
         answer = options["yes"]
+        truncate = options["truncate"]
         before = options["before_date"]
-
+        if truncate and before:
+            self.stdout.write(
+                "Truncate deletes all log entries and can not be passed with before-date."
+            )
+            return
         if answer is None:
             warning_message = (
                 "This action will clear all log entries from the database."
@@ -42,11 +56,39 @@ def handle(self, *args, **options):
             )
             answer = response == "y"
 
-        if answer:
+        if not answer:
+            self.stdout.write("Aborted.")
+            return
+
+        if not truncate:
             entries = LogEntry.objects.all()
             if before is not None:
                 entries = entries.filter(timestamp__date__lt=before)
             count, _ = entries.delete()
             self.stdout.write("Deleted %d objects." % count)
         else:
-            self.stdout.write("Aborted.")
+            database_vendor = connection.vendor
+            database_display_name = connection.display_name
+            table_name = LogEntry._meta.db_table
+            if not TruncateQuery.support_truncate_statement(database_vendor):
+                self.stdout.write(
+                    "Database %s does not support truncate statement."
+                    % database_display_name
+                )
+                return
+            with connection.cursor() as cursor:
+                query = TruncateQuery.to_sql(table_name)
+                cursor.execute(query)
+            self.stdout.write("Truncated log entry table.")
+
+
+class TruncateQuery:
+    SUPPORTED_VENDORS = ("postgresql", "mysql", "sqlite", "oracle", "microsoft")
+
+    @classmethod
+    def support_truncate_statement(cls, database_vendor) -> bool:
+        return database_vendor in cls.SUPPORTED_VENDORS
+
+    @staticmethod
+    def to_sql(table_name) -> str:
+        return f"TRUNCATE TABLE {table_name};"

auditlog_tests/test_commands.py

@@ -6,7 +6,7 @@
 
 import freezegun
 from django.core.management import call_command
-from django.test import TestCase
+from django.test import TestCase, TransactionTestCase
 
 from auditlog_tests.models import SimpleModel
 
@@ -110,3 +110,82 @@ def test_before_date(self):
             out, "Deleted 1 objects.", msg="Output shows deleted 1 object."
         )
         self.assertEqual(err, "", msg="No stderr")
+
+
+class AuditlogFlushWithTruncateTest(TransactionTestCase):
+    def setUp(self):
+        input_patcher = mock.patch("builtins.input")
+        self.mock_input = input_patcher.start()
+        self.addCleanup(input_patcher.stop)
+
+    def make_object(self):
+        return SimpleModel.objects.create(text="I am a simple model.")
+
+    def call_command(self, *args, **kwargs):
+        outbuf = StringIO()
+        errbuf = StringIO()
+        call_command("auditlogflush", *args, stdout=outbuf, stderr=errbuf, **kwargs)
+        return outbuf.getvalue().strip(), errbuf.getvalue().strip()
+
+    def test_flush_with_both_truncate_and_before_date_options(self):
+        obj = self.make_object()
+        self.assertEqual(obj.history.count(), 1, msg="There is one log entry.")
+        out, err = self.call_command("--truncate", "--before-date=2000-01-01")
+
+        self.assertEqual(obj.history.count(), 1, msg="There is still one log entry.")
+        self.assertEqual(
+            out,
+            "Truncate deletes all log entries and can not be passed with before-date.",
+            msg="Output shows error",
+        )
+        self.assertEqual(err, "", msg="No stderr")
+
+    def test_flush_with_truncate_and_yes(self):
+        obj = self.make_object()
+        self.assertEqual(obj.history.count(), 1, msg="There is one log entry.")
+        out, err = self.call_command("--truncate", "--y")
+
+        self.assertEqual(obj.history.count(), 0, msg="There is no log entry.")
+        self.assertEqual(
+            out,
+            "Truncated log entry table.",
+            msg="Output shows table gets truncate",
+        )
+        self.assertEqual(err, "", msg="No stderr")
+
+    def test_flush_with_truncate_with_input_yes(self):
+        obj = self.make_object()
+        self.assertEqual(obj.history.count(), 1, msg="There is one log entry.")
+        self.mock_input.return_value = "Y\n"
+        out, err = self.call_command("--truncate")
+
+        self.assertEqual(obj.history.count(), 0, msg="There is no log entry.")
+        self.assertEqual(
+            out,
+            "This action will clear all log entries from the database.\nTruncated log entry table.",
+            msg="Output shows warning and table gets truncate",
+        )
+        self.assertEqual(err, "", msg="No stderr")
+
+    @mock.patch(
+        "django.db.connection.vendor",
+        new_callable=mock.PropertyMock(return_value="unknown"),
+    )
+    @mock.patch(
+        "django.db.connection.display_name",
+        new_callable=mock.PropertyMock(return_value="Unknown"),
+    )
+    def test_flush_with_truncate_for_unsupported_database_vendor(
+        self, mocked_vendor, mocked_db_name
+    ):
+        obj = self.make_object()
+        self.assertEqual(obj.history.count(), 1, msg="There is one log entry.")
+        out, err = self.call_command("--truncate", "--y")
+
+        self.assertEqual(obj.history.count(), 1, msg="There is still one log entry.")
+        self.assertEqual(
+            out,
+            "Database Unknown does not support truncate statement.",
+            msg="Output shows error",
+        )
+        self.assertEqual(err, "", msg="No stderr")