optimize flush audit log

Author: valentijnscholten

dojo/management/commands/flush_auditlog.py

@@ -0,0 +1,27 @@
+from django.core.management.base import BaseCommand
+
+from dojo.tasks import run_flush_auditlog
+
+
+class Command(BaseCommand):
+    help = "Flush old audit log entries based on retention and batching settings"
+
+    def add_arguments(self, parser):
+        parser.add_argument("--retention-months", type=int, default=None, help="Override retention period in months")
+        parser.add_argument("--batch-size", type=int, default=None, help="Override batch size")
+        parser.add_argument("--max-batches", type=int, default=None, help="Override max batches per run")
+
+    def handle(self, *args, **options):
+        deleted_total, batches_done, reached_limit = run_flush_auditlog(
+            retention_period=options.get("retention_months"),
+            batch_size=options.get("batch_size"),
+            max_batches=options.get("max_batches"),
+        )
+        if reached_limit:
+            self.stdout.write(self.style.WARNING(
+                f"Reached max batches limit; deleted {deleted_total} entries in {batches_done} batches.",
+            ))
+        else:
+            self.stdout.write(self.style.SUCCESS(
+                f"Deleted {deleted_total} audit log entries in {batches_done} batches.",
+            ))

dojo/settings/settings.dist.py

@@ -260,6 +260,10 @@
     DD_TRACK_IMPORT_HISTORY=(bool, True),
     # Delete Auditlogs older than x month; -1 to keep all logs
     DD_AUDITLOG_FLUSH_RETENTION_PERIOD=(int, -1),
+    # Batch size for flushing audit logs per task run
+    DD_AUDITLOG_FLUSH_BATCH_SIZE=(int, 1000),
+    # Maximum number of batches to process per task run
+    DD_AUDITLOG_FLUSH_MAX_BATCHES=(int, 100),
     # Allow grouping of findings in the same test, for example to group findings per dependency
     # DD_FEATURE_FINDING_GROUPS feature is moved to system_settings, will be removed from settings file
     DD_FEATURE_FINDING_GROUPS=(bool, True),
@@ -1183,7 +1187,7 @@ def saml2_attrib_map_format(din):
     },
     "flush_auditlog": {
         "task": "dojo.tasks.flush_auditlog",
-        "schedule": timedelta(hours=8),
+        "schedule": timedelta(minutes=1),
     },
     "update-findings-from-source-issues": {
         "task": "dojo.tools.tool_issue_updater.update_findings_from_source_issues",
@@ -1914,6 +1918,8 @@ def saml2_attrib_map_format(din):
 AUDITLOG_TYPE = env("DD_AUDITLOG_TYPE")
 AUDITLOG_TWO_STEP_MIGRATION = False
 AUDITLOG_USE_TEXT_CHANGES_IF_JSON_IS_NOT_PRESENT = False
+AUDITLOG_FLUSH_BATCH_SIZE = env("DD_AUDITLOG_FLUSH_BATCH_SIZE")
+AUDITLOG_FLUSH_MAX_BATCHES = env("DD_AUDITLOG_FLUSH_MAX_BATCHES")
 
 USE_FIRST_SEEN = env("DD_USE_FIRST_SEEN")
 USE_QUALYS_LEGACY_SEVERITY_PARSING = env("DD_QUALYS_LEGACY_SEVERITY_PARSING")

dojo/tasks.py

@@ -1,5 +1,5 @@
 import logging
-from datetime import date, timedelta
+from datetime import date, datetime, time, timedelta
 
 from auditlog.models import LogEntry
 from celery.utils.log import get_task_logger
@@ -93,24 +93,57 @@ def cleanup_alerts(*args, **kwargs):
         logger.info("total number of alerts deleted: %s", total_deleted_count)
 
 
-@app.task(bind=True)
-def flush_auditlog(*args, **kwargs):
-    retention_period = settings.AUDITLOG_FLUSH_RETENTION_PERIOD
+def run_flush_auditlog(retention_period: int | None = None,
+                       batch_size: int | None = None,
+                       max_batches: int | None = None) -> tuple[int, int, bool]:
+    """
+    Deletes audit log entries older than the configured retention period.
 
+    Returns a tuple of (deleted_total, batches_done, reached_limit).
+    """
+    retention_period = retention_period if retention_period is not None else getattr(settings, "AUDITLOG_FLUSH_RETENTION_PERIOD", -1)
     if retention_period < 0:
         logger.info("Flushing auditlog is disabled")
-        return
+        return 0, 0, False
 
     logger.info("Running Cleanup Task for Logentries with %d Months retention", retention_period)
-    retention_date = date.today() - relativedelta(months=retention_period)
-    subset = LogEntry.objects.filter(timestamp__date__lt=retention_date)
-    event_count = subset.count()
-    logger.debug("Initially received %d Logentries", event_count)
-    if event_count > 0:
-        subset._raw_delete(subset.db)
-        logger.debug("Total number of audit log entries deleted: %s", event_count)
+    # Compute a datetime cutoff at start of the cutoff day to keep index-usage friendly
+    retention_day = date.today() - relativedelta(months=retention_period)
+    # Use a timestamp to avoid postgres having to cast to a Date field
+    cutoff_dt = datetime.combine(retention_day, time.min, tzinfo=timezone.get_current_timezone())
+
+    # Settings to control batching; sensible defaults if not configured
+    batch_size = batch_size if batch_size is not None else getattr(settings, "AUDITLOG_FLUSH_BATCH_SIZE", 1000)
+    max_batches = max_batches if max_batches is not None else getattr(settings, "AUDITLOG_FLUSH_MAX_BATCHES", 100)
+
+    # Delete in batches to avoid long-running transactions and table locks
+    deleted_total = 0
+    batches_done = 0
+    while batches_done < max_batches:
+        batch_qs = LogEntry.objects.filter(timestamp__lt=cutoff_dt).order_by("pk")
+        pks = list(batch_qs.values_list("pk", flat=True)[:batch_size])
+        if not pks:
+            if batches_done == 0:
+                logger.info("No outdated Logentries found")
+            break
+        qs = LogEntry.objects.filter(pk__in=pks)
+        deleted_count = qs._raw_delete(qs.db)
+        deleted_total += int(deleted_count)
+        batches_done += 1
+        logger.info("Deleted batch %s (size ~%s), total deleted: %s", batches_done, batch_size, deleted_total)
+
+    reached_limit = batches_done >= max_batches
+    if reached_limit:
+        logger.info("Reached max batches limit (%s). Remaining audit log entries will be deleted in the next run.", max_batches)
     else:
-        logger.debug("No outdated Logentries found")
+        logger.info("Total number of audit log entries deleted: %s", deleted_total)
+
+    return deleted_total, batches_done, reached_limit
+
+
+@app.task(bind=True)
+def flush_auditlog(*args, **kwargs):
+    run_flush_auditlog()
 
 
 @app.task(bind=True)