from django.core.management.base import BaseCommand from worktimeservice.tasks import catch_up_work_balances_for_user, schedule_catch_up_balances, normalize_user_id from stamps.models import Stamp import json class Command(BaseCommand): help = 'Manually trigger balance catch-up task. This will recalculate balances for users, including dates with 0.0 balances that have stamps.' def add_arguments(self, parser): parser.add_argument( '--user-id', type=str, help='Specific user ID to recalculate balances for (if not provided, processes all users)', ) parser.add_argument( '--async', action='store_true', help='Run asynchronously via Celery (default: synchronous execution)', ) parser.add_argument( '--all-users', action='store_true', help='Process all users with stamps (uses schedule_catch_up_balances)', ) parser.add_argument( '--verbose', action='store_true', help='Show detailed output', ) def handle(self, *args, **options): user_id_raw = options.get('user_id', None) is_async = options.get('async', False) all_users = options.get('all_users', False) verbose = options.get('verbose', False) # Normalize user_id if provided if user_id_raw: user_id = normalize_user_id(user_id_raw) if user_id != user_id_raw: self.stdout.write(self.style.WARNING(f"User ID normalized: '{user_id_raw}' -> '{user_id}'")) else: user_id = None if all_users: # Process all users if is_async: self.stdout.write(self.style.WARNING('Queuing balance catch-up for all users asynchronously...')) self.stdout.write('Note: Celery worker must be running for this to execute.') try: task = schedule_catch_up_balances.delay() self.stdout.write('') self.stdout.write(self.style.SUCCESS('✓ Task queued successfully!')) self.stdout.write(f'Task ID: {task.id}') self.stdout.write('') self.stdout.write('Check Celery logs to see progress.') except Exception as e: self.stdout.write('') self.stdout.write(self.style.ERROR(f'✗ Failed to queue task: {str(e)}')) if verbose: import traceback self.stdout.write('') self.stdout.write('Traceback:') self.stdout.write(traceback.format_exc()) else: self.stdout.write(self.style.WARNING('Running balance catch-up for all users synchronously...')) self.stdout.write('This may take a while depending on the number of users.') self.stdout.write('') try: result = schedule_catch_up_balances() if verbose: self.stdout.write('') self.stdout.write(self.style.SUCCESS('Task completed successfully!')) self.stdout.write('') self.stdout.write('Result:') self.stdout.write(json.dumps(result, indent=2, default=str)) else: self.stdout.write(self.style.SUCCESS('✓ Balance catch-up completed for all users')) except Exception as e: self.stdout.write('') self.stdout.write(self.style.ERROR(f'✗ Failed: {str(e)}')) if verbose: import traceback self.stdout.write('') self.stdout.write('Traceback:') self.stdout.write(traceback.format_exc()) elif user_id: # Process specific user if is_async: self.stdout.write(self.style.WARNING(f'Queuing balance catch-up for user {user_id} asynchronously...')) self.stdout.write('Note: Celery worker must be running for this to execute.') try: task = catch_up_work_balances_for_user.delay(user_id) self.stdout.write('') self.stdout.write(self.style.SUCCESS('✓ Task queued successfully!')) self.stdout.write(f'Task ID: {task.id}') self.stdout.write('') self.stdout.write('Check Celery logs to see progress.') except Exception as e: self.stdout.write('') self.stdout.write(self.style.ERROR(f'✗ Failed to queue task: {str(e)}')) if verbose: import traceback self.stdout.write('') self.stdout.write('Traceback:') self.stdout.write(traceback.format_exc()) else: self.stdout.write(self.style.WARNING(f'Running balance catch-up for user {user_id} synchronously...')) self.stdout.write('This will recalculate balances for the last 30 days and dates with 0.0 balances.') self.stdout.write('') try: # Call the task synchronously using .apply() which handles the self parameter correctly # .apply() runs the task synchronously in the current process if verbose: self.stdout.write(f'Calling catch_up_work_balances_for_user for user: {user_id}') # Use .apply() to run synchronously - this properly handles bind=True tasks # .apply() returns an AsyncResult even when running synchronously task_result = catch_up_work_balances_for_user.apply(args=[user_id]) # Check if the task failed if task_result.failed(): error = task_result.result self.stdout.write(self.style.ERROR(f'✗ Task failed: {error}')) if verbose: import traceback if isinstance(error, Exception): self.stdout.write(traceback.format_exception(type(error), error, error.__traceback__)) return # Extract the actual result from the AsyncResult result = task_result.result if verbose: self.stdout.write(f'Function returned result with status: {result.get("status") if result else "None"}') if verbose: self.stdout.write('') self.stdout.write(self.style.SUCCESS('Task completed successfully!')) self.stdout.write('') self.stdout.write('Result:') self.stdout.write(json.dumps(result, indent=2, default=str)) # Show errors if any if result.get('errors'): self.stdout.write('') self.stdout.write(self.style.WARNING('Errors encountered:')) for error in result.get('errors', []): self.stdout.write(f' - {error}') else: status = result.get('status', 'unknown') if status == 'ok': dates_processed = result.get('dates_processed', 0) total_checked = result.get('total_dates_checked', 0) self.stdout.write('') self.stdout.write(self.style.SUCCESS(f'✓ Balance catch-up completed!')) self.stdout.write(f'Dates processed: {dates_processed} out of {total_checked} checked') if result.get('errors'): self.stdout.write(self.style.WARNING(f'⚠ {len(result.get("errors", []))} errors occurred (use --verbose to see details)')) elif status == 'no_data': self.stdout.write('') self.stdout.write(self.style.WARNING('No stamps found for this user')) else: self.stdout.write('') self.stdout.write(self.style.ERROR(f'Task failed with status: {status}')) if 'error' in result: self.stdout.write(f'Error: {result.get("error")}') except Exception as e: self.stdout.write('') self.stdout.write(self.style.ERROR(f'✗ Failed: {str(e)}')) if verbose: import traceback self.stdout.write('') self.stdout.write('Traceback:') self.stdout.write(traceback.format_exc()) else: self.stdout.write(self.style.ERROR('Please specify --user-id USER_ID or --all-users')) self.stdout.write('') self.stdout.write('Examples:') self.stdout.write(' python manage.py catch_up_balances --user-id edd631c63c1848c681244f5c7c868801') self.stdout.write(' python manage.py catch_up_balances --all-users') self.stdout.write(' python manage.py catch_up_balances --user-id USER_ID --async') self.stdout.write(' python manage.py catch_up_balances --all-users --async')