# -*- coding: utf-8 -*- """ Cleanup Notifications Management Command This command cleans up old, expired, and unnecessary notifications to maintain database performance and storage efficiency. Usage: python manage.py cleanup_notifications python manage.py cleanup_notifications --days 30 python manage.py cleanup_notifications --expired-only python manage.py cleanup_notifications --dry-run Author: AdTlas Development Team Version: 1.0.0 Last Updated: 2024 """ import logging from datetime import timedelta from django.core.management.base import BaseCommand, CommandError from django.utils import timezone from django.db import transaction from django.db.models import Count, Q from apps.notifications.models import ( Notification, NotificationHistory, NotificationQueue, ) logger = logging.getLogger(__name__) class Command(BaseCommand): """ Management command to clean up notifications. This command provides various cleanup options for maintaining the notification system's database tables. """ help = 'Clean up old and expired notifications' def add_arguments(self, parser): """ Add command line arguments. Args: parser: Argument parser instance """ # Cleanup scope options parser.add_argument( '--days', type=int, default=30, help='Number of days to keep notifications (default: 30)' ) parser.add_argument( '--expired-only', action='store_true', help='Only clean up expired notifications' ) parser.add_argument( '--read-only', action='store_true', help='Only clean up read notifications' ) parser.add_argument( '--archived-only', action='store_true', help='Only clean up archived notifications' ) # Cleanup targets parser.add_argument( '--notifications', action='store_true', default=True, help='Clean up notifications (default: True)' ) parser.add_argument( '--history', action='store_true', help='Clean up notification history' ) parser.add_argument( '--queue', action='store_true', help='Clean up notification queue' ) parser.add_argument( '--all', action='store_true', help='Clean up all notification data' ) # Processing options parser.add_argument( '--batch-size', type=int, default=1000, help='Batch size for deletion (default: 1000)' ) parser.add_argument( '--dry-run', action='store_true', help='Show what would be deleted without actually deleting' ) parser.add_argument( '--force', action='store_true', help='Skip confirmation prompts' ) parser.add_argument( '--verbose', action='store_true', help='Show detailed progress information' ) def handle(self, *args, **options): """ Handle the command execution. Args: *args: Positional arguments **options: Command options """ try: # Validate options self._validate_options(options) # Calculate cutoff date cutoff_date = timezone.now() - timedelta(days=options['days']) # Show summary self._show_summary(options, cutoff_date) # Get confirmation if not forced if not options['force'] and not options['dry_run']: if not self._get_confirmation(): self.stdout.write("Cleanup cancelled.") return # Perform cleanup results = self._perform_cleanup(options, cutoff_date) # Show results self._show_results(results, options) except Exception as e: logger.error(f"Error in cleanup_notifications command: {e}") raise CommandError(f"Command failed: {e}") def _validate_options(self, options): """ Validate command options. Args: options: Command options Raises: CommandError: If options are invalid """ if options['days'] < 1: raise CommandError("Days must be at least 1") if options['batch_size'] < 1: raise CommandError("Batch size must be at least 1") # Set cleanup targets if --all is specified if options['all']: options['notifications'] = True options['history'] = True options['queue'] = True def _show_summary(self, options, cutoff_date): """ Show cleanup summary. Args: options: Command options cutoff_date: Cutoff date for cleanup """ self.stdout.write(self.style.SUCCESS("\n=== Notification Cleanup Summary ===")) self.stdout.write(f"Cutoff Date: {cutoff_date.strftime('%Y-%m-%d %H:%M:%S')}") self.stdout.write(f"Days to Keep: {options['days']}") self.stdout.write(f"Batch Size: {options['batch_size']}") # Show scope scope = [] if options['expired_only']: scope.append("expired only") if options['read_only']: scope.append("read only") if options['archived_only']: scope.append("archived only") if scope: self.stdout.write(f"Scope: {', '.join(scope)}") else: self.stdout.write("Scope: all notifications") # Show targets targets = [] if options['notifications']: targets.append("notifications") if options['history']: targets.append("history") if options['queue']: targets.append("queue") self.stdout.write(f"Targets: {', '.join(targets)}") if options['dry_run']: self.stdout.write(self.style.WARNING("DRY RUN: No data will be deleted")) # Show counts self._show_current_counts(options, cutoff_date) self.stdout.write("\n" + "="*50) def _show_current_counts(self, options, cutoff_date): """ Show current counts of data to be cleaned. Args: options: Command options cutoff_date: Cutoff date for cleanup """ self.stdout.write("\nCurrent Counts:") if options['notifications']: notification_count = self._get_notification_count(options, cutoff_date) self.stdout.write(f" Notifications to clean: {notification_count:,}") if options['history']: history_count = self._get_history_count(options, cutoff_date) self.stdout.write(f" History records to clean: {history_count:,}") if options['queue']: queue_count = self._get_queue_count(options, cutoff_date) self.stdout.write(f" Queue items to clean: {queue_count:,}") def _get_notification_count(self, options, cutoff_date): """ Get count of notifications to be cleaned. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: int: Count of notifications """ queryset = Notification.objects.all() if options['expired_only']: queryset = queryset.filter( expires_at__lt=timezone.now() ) elif options['read_only']: queryset = queryset.filter( is_read=True, created_at__lt=cutoff_date ) elif options['archived_only']: queryset = queryset.filter( is_archived=True, created_at__lt=cutoff_date ) else: queryset = queryset.filter( Q(created_at__lt=cutoff_date) | Q(expires_at__lt=timezone.now()) ) return queryset.count() def _get_history_count(self, options, cutoff_date): """ Get count of history records to be cleaned. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: int: Count of history records """ return NotificationHistory.objects.filter( attempted_at__lt=cutoff_date ).count() def _get_queue_count(self, options, cutoff_date): """ Get count of queue items to be cleaned. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: int: Count of queue items """ return NotificationQueue.objects.filter( created_at__lt=cutoff_date, status__in=['completed', 'failed'] ).count() def _get_confirmation(self): """ Get user confirmation for cleanup. Returns: bool: True if user confirms """ response = input("\nProceed with cleanup? [y/N]: ") return response.lower() in ['y', 'yes'] def _perform_cleanup(self, options, cutoff_date): """ Perform the actual cleanup. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: dict: Cleanup results """ results = { 'notifications_deleted': 0, 'history_deleted': 0, 'queue_deleted': 0, 'total_deleted': 0, 'errors': [] } try: with transaction.atomic(): if options['notifications']: count = self._cleanup_notifications(options, cutoff_date) results['notifications_deleted'] = count results['total_deleted'] += count if options['history']: count = self._cleanup_history(options, cutoff_date) results['history_deleted'] = count results['total_deleted'] += count if options['queue']: count = self._cleanup_queue(options, cutoff_date) results['queue_deleted'] = count results['total_deleted'] += count if options['dry_run']: # Rollback transaction for dry run transaction.set_rollback(True) except Exception as e: logger.error(f"Error during cleanup: {e}") results['errors'].append(str(e)) return results def _cleanup_notifications(self, options, cutoff_date): """ Clean up notifications. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: int: Number of notifications deleted """ self.stdout.write("\nCleaning up notifications...") # Build queryset queryset = Notification.objects.all() if options['expired_only']: queryset = queryset.filter( expires_at__lt=timezone.now() ) elif options['read_only']: queryset = queryset.filter( is_read=True, created_at__lt=cutoff_date ) elif options['archived_only']: queryset = queryset.filter( is_archived=True, created_at__lt=cutoff_date ) else: queryset = queryset.filter( Q(created_at__lt=cutoff_date) | Q(expires_at__lt=timezone.now()) ) return self._delete_in_batches( queryset, options['batch_size'], "notifications", options['verbose'] ) def _cleanup_history(self, options, cutoff_date): """ Clean up notification history. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: int: Number of history records deleted """ self.stdout.write("\nCleaning up notification history...") queryset = NotificationHistory.objects.filter( attempted_at__lt=cutoff_date ) return self._delete_in_batches( queryset, options['batch_size'], "history records", options['verbose'] ) def _cleanup_queue(self, options, cutoff_date): """ Clean up notification queue. Args: options: Command options cutoff_date: Cutoff date for cleanup Returns: int: Number of queue items deleted """ self.stdout.write("\nCleaning up notification queue...") queryset = NotificationQueue.objects.filter( created_at__lt=cutoff_date, status__in=['completed', 'failed'] ) return self._delete_in_batches( queryset, options['batch_size'], "queue items", options['verbose'] ) def _delete_in_batches(self, queryset, batch_size, item_type, verbose): """ Delete records in batches. Args: queryset: QuerySet to delete from batch_size: Size of each batch item_type: Type of items being deleted (for logging) verbose: Whether to show verbose output Returns: int: Total number of records deleted """ total_deleted = 0 total_count = queryset.count() if total_count == 0: if verbose: self.stdout.write(f" No {item_type} to delete") return 0 if verbose: self.stdout.write(f" Deleting {total_count:,} {item_type} in batches of {batch_size:,}") while True: # Get batch of IDs batch_ids = list(queryset.values_list('id', flat=True)[:batch_size]) if not batch_ids: break # Delete batch deleted_count = queryset.filter(id__in=batch_ids).delete()[0] total_deleted += deleted_count if verbose: progress = (total_deleted / total_count) * 100 self.stdout.write( f"\r Progress: {total_deleted:,}/{total_count:,} ({progress:.1f}%)", ending='' ) if verbose: self.stdout.write(f"\n Deleted {total_deleted:,} {item_type}") return total_deleted def _show_results(self, results, options): """ Show cleanup results. Args: results: Cleanup results options: Command options """ self.stdout.write("\n=== Cleanup Results ===") if options['dry_run']: self.stdout.write(self.style.WARNING("DRY RUN - No data was actually deleted")) if results['notifications_deleted'] > 0: self.stdout.write( self.style.SUCCESS(f"✓ Notifications deleted: {results['notifications_deleted']:,}") ) if results['history_deleted'] > 0: self.stdout.write( self.style.SUCCESS(f"✓ History records deleted: {results['history_deleted']:,}") ) if results['queue_deleted'] > 0: self.stdout.write( self.style.SUCCESS(f"✓ Queue items deleted: {results['queue_deleted']:,}") ) if results['total_deleted'] > 0: self.stdout.write( self.style.SUCCESS(f"\n✓ Total records deleted: {results['total_deleted']:,}") ) else: self.stdout.write("No records were deleted.") if results['errors']: self.stdout.write("\nErrors:") for error in results['errors']: self.stdout.write(self.style.ERROR(f"✗ {error}")) # Show remaining counts self._show_remaining_counts() self.stdout.write("\nCleanup completed!") def _show_remaining_counts(self): """ Show remaining counts after cleanup. """ try: notification_count = Notification.objects.count() history_count = NotificationHistory.objects.count() queue_count = NotificationQueue.objects.count() self.stdout.write("\nRemaining Records:") self.stdout.write(f" Notifications: {notification_count:,}") self.stdout.write(f" History: {history_count:,}") self.stdout.write(f" Queue: {queue_count:,}") except Exception as e: logger.error(f"Error getting remaining counts: {e}")