pghistory_backfill: avoid prefetching - dry-run working

Author: valentijnscholten

dojo/management/commands/pghistory_backfill.py

@@ -4,6 +4,7 @@
 This command creates initial snapshots for all existing records in tracked models.
 """
 import logging
+import time
 
 from django.apps import apps
 from django.conf import settings
@@ -33,6 +34,16 @@ def add_arguments(self, parser):
             action="store_true",
             help="Show what would be done without actually creating events",
         )
+        parser.add_argument(
+            "--log-queries",
+            action="store_true",
+            help="Enable database query logging (default: enabled)",
+        )
+        parser.add_argument(
+            "--no-log-queries",
+            action="store_true",
+            help="Disable database query logging",
+        )
 
     def get_excluded_fields(self, model_name):
         """Get the list of excluded fields for a specific model from pghistory configuration."""
@@ -45,6 +56,51 @@ def get_excluded_fields(self, model_name):
         }
         return excluded_fields_map.get(model_name, [])
 
+    def enable_db_logging(self):
+        """Enable database query logging for this command."""
+        # Store original DEBUG setting
+        self.original_debug = settings.DEBUG
+
+        # Configure database query logging
+        db_logger = logging.getLogger("django.db.backends")
+        db_logger.setLevel(logging.DEBUG)
+
+        # Add a handler if one doesn't exist
+        if not db_logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+            )
+            handler.setFormatter(formatter)
+            db_logger.addHandler(handler)
+
+        # Also enable the SQL logger specifically
+        sql_logger = logging.getLogger("django.db.backends.sql")
+        sql_logger.setLevel(logging.DEBUG)
+
+        # Ensure the root logger propagates to our handlers
+        if not sql_logger.handlers:
+            sql_logger.addHandler(handler)
+
+        # Enable query logging in Django settings
+        settings.DEBUG = True
+
+        self.stdout.write(
+            self.style.SUCCESS("Database query logging enabled"),
+        )
+
+    def disable_db_logging(self):
+        """Disable database query logging."""
+        # Restore original DEBUG setting
+        settings.DEBUG = self.original_debug
+
+        # Disable query logging by setting a higher level
+        logging.getLogger("django.db.backends").setLevel(logging.INFO)
+        logging.getLogger("django.db.backends.sql").setLevel(logging.INFO)
+        self.stdout.write(
+            self.style.SUCCESS("Database query logging disabled"),
+        )
+
     def handle(self, *args, **options):
         if not settings.ENABLE_AUDITLOG or settings.AUDITLOG_TYPE != "django-pghistory":
             self.stdout.write(
@@ -55,6 +111,17 @@ def handle(self, *args, **options):
             )
             return
 
+        # Enable database query logging based on options
+        # Default to enabled unless explicitly disabled
+        enable_query_logging = not options.get("no_log_queries")
+
+        if enable_query_logging:
+            self.enable_db_logging()
+        else:
+            self.stdout.write(
+                self.style.WARNING("Database query logging disabled"),
+            )
+
         # Models that are tracked by pghistory
         tracked_models = [
             "Dojo_User", "Endpoint", "Engagement", "Finding", "Finding_Group",
@@ -83,9 +150,11 @@ def handle(self, *args, **options):
             )
 
         total_processed = 0
+        total_start_time = time.time()
         self.stdout.write(f"Starting backfill for {len(tracked_models)} model(s)...")
 
         for model_name in tracked_models:
+            model_start_time = time.time()
             self.stdout.write(f"\nProcessing {model_name}...")
 
             try:
@@ -143,6 +212,7 @@ def handle(self, *args, **options):
                 processed = 0
                 event_records = []
                 failed_records = []
+                batch_start_time = time.time()
 
                 for instance in records_needing_backfill.iterator():
                     try:
@@ -156,8 +226,17 @@ def handle(self, *args, **options):
                         for field in instance._meta.fields:
                             field_name = field.name
                             if field_name not in excluded_fields:
-                                field_value = getattr(instance, field_name)
-                                event_data[field_name] = field_value
+                                # Handle foreign key fields differently
+                                if field.many_to_one:  # ForeignKey field
+                                    # For foreign keys, use the _id field to get the raw ID value
+                                    # Store it under the _id field name for the Event model
+                                    field_id_name = f"{field_name}_id"
+                                    field_value = getattr(instance, field_id_name)
+                                    event_data[field_id_name] = field_value
+                                else:
+                                    # For non-foreign key fields, use value_from_object() to avoid queries
+                                    field_value = field.value_from_object(instance)
+                                    event_data[field_name] = field_value
 
                         # Explicitly preserve created timestamp from the original instance
                         # Only if not excluded and exists
@@ -180,12 +259,16 @@ def handle(self, *args, **options):
 
                     except Exception as e:
                         failed_records.append(instance.id)
-                        logger.error(
-                            f"Failed to prepare event for {model_name} ID {instance.id}: {e}",
+                        logger.exception(
+                            f"Failed to prepare event for {model_name} ID {instance.id}",
                         )
 
                     # Bulk create when we hit batch_size records
                     if len(event_records) >= batch_size:
+                        batch_end_time = time.time()
+                        batch_duration = batch_end_time - batch_start_time
+                        batch_records_per_second = len(event_records) / batch_duration if batch_duration > 0 else 0
+
                         if not dry_run and event_records:
                             try:
                                 attempted = len(event_records)
@@ -199,19 +282,25 @@ def handle(self, *args, **options):
                                         f"actually created {actually_created} ({attempted - actually_created} skipped)",
                                     )
                             except Exception as e:
-                                logger.error(f"Failed to bulk create events for {model_name}: {e}")
+                                logger.exception(f"Failed to bulk create events for {model_name}")
                                 raise
                         elif dry_run:
                             processed += len(event_records)
 
                         event_records = []  # Reset for next batch
+                        batch_start_time = time.time()  # Reset batch timer
 
-                        # Progress update
+                        # Progress update with batch timing
                         progress = (processed / backfill_count) * 100
-                        self.stdout.write(f"  Processed {processed:,}/{backfill_count:,} records needing backfill ({progress:.1f}%)")
+                        self.stdout.write(f"  Processed {processed:,}/{backfill_count:,} records needing backfill ({progress:.1f}%) - "
+                                        f"Last batch: {batch_duration:.2f}s ({batch_records_per_second:.1f} records/sec)")
 
                 # Handle remaining records
                 if event_records:
+                    batch_end_time = time.time()
+                    batch_duration = batch_end_time - batch_start_time
+                    batch_records_per_second = len(event_records) / batch_duration if batch_duration > 0 else 0
+
                     if not dry_run:
                         try:
                             attempted = len(event_records)
@@ -225,41 +314,61 @@ def handle(self, *args, **options):
                                     f"actually created {actually_created} ({attempted - actually_created} skipped)",
                                 )
                         except Exception as e:
-                            logger.error(f"Failed to bulk create final batch for {model_name}: {e}")
+                            logger.exception(f"Failed to bulk create final batch for {model_name}")
                             raise
                     else:
                         processed += len(event_records)
 
+                    # Log final batch timing
+                    self.stdout.write(f"  Final batch: {batch_duration:.2f}s ({batch_records_per_second:.1f} records/sec)")
+
                 # Final progress update
                 if backfill_count > 0:
                     progress = (processed / backfill_count) * 100
                     self.stdout.write(f"  Processed {processed:,}/{backfill_count:,} records needing backfill ({progress:.1f}%)")
 
                 total_processed += processed
 
-                # Show completion summary
+                # Calculate timing for this model
+                model_end_time = time.time()
+                model_duration = model_end_time - model_start_time
+                records_per_second = processed / model_duration if model_duration > 0 else 0
+
+                # Show completion summary with timing
                 if failed_records:
                     self.stdout.write(
                         self.style.WARNING(
                             f"  ⚠ Completed {model_name}: {processed:,} records processed, "
-                            f"{len(failed_records)} records failed",
+                            f"{len(failed_records)} records failed in {model_duration:.2f}s "
+                            f"({records_per_second:.1f} records/sec)",
                         ),
                     )
                 else:
                     self.stdout.write(
                         self.style.SUCCESS(
-                            f"  ✓ Completed {model_name}: {processed:,} records",
+                            f"  ✓ Completed {model_name}: {processed:,} records in {model_duration:.2f}s "
+                            f"({records_per_second:.1f} records/sec)",
                         ),
                     )
 
             except Exception as e:
                 self.stdout.write(
                     self.style.ERROR(f"  ✗ Failed to process {model_name}: {e}"),
                 )
-                logger.error(f"Error processing {model_name}: {e}")
+                logger.exception(f"Error processing {model_name}")
+
+        # Calculate total timing
+        total_end_time = time.time()
+        total_duration = total_end_time - total_start_time
+        total_records_per_second = total_processed / total_duration if total_duration > 0 else 0
+
+        # Disable database query logging if it was enabled
+        if enable_query_logging:
+            self.disable_db_logging()
 
         self.stdout.write(
             self.style.SUCCESS(
-                f"\nBACKFILL COMPLETE: Processed {total_processed:,} records",
+                f"\nBACKFILL COMPLETE: Processed {total_processed:,} records in {total_duration:.2f}s "
+                f"({total_records_per_second:.1f} records/sec)",
             ),
         )