diff --git a/api/views.py b/api/views.py index 35aaa7e..92f0ca8 100644 --- a/api/views.py +++ b/api/views.py @@ -23,7 +23,7 @@ from django.views.decorators.http import require_http_methods from api.models import WireguardStatusCache from cluster.models import ClusterSettings, WorkerStatus -from scheduler.models import PeerScheduling +from scheduler.models import PeerScheduling, ScheduleProfile from user_manager.models import AuthenticationToken, UserAcl from vpn_invite.models import InviteSettings, PeerInvite from wgwadmlibrary.tools import create_peer_invite, get_peer_invite_data, send_email, user_allowed_peers, \ @@ -413,6 +413,76 @@ def cron_refresh_wireguard_status_cache(request): return JsonResponse(data) +def cron_calculate_peer_schedules(request): + data = { + 'status': 'success', + 'updated_records': 0, + 'skipped_records': 0, + } + + peer_scheduling_queryset = ( + PeerScheduling.objects + .select_related('peer', 'profile') + .filter( + profile__active=True, + peer__suspended=False, + ) + .filter( + Q(next_scheduled_enable_at__isnull=True) | + Q(next_scheduled_disable_at__isnull=True) + ) + ) + + if not peer_scheduling_queryset.exists(): + return JsonResponse(data) + + distinct_profile_ids = list( + peer_scheduling_queryset.values_list('profile_id', flat=True).distinct() + ) + + schedule_profiles = ( + ScheduleProfile.objects + .filter(id__in=distinct_profile_ids, active=True) + .prefetch_related('time_interval') + ) + + profile_next_dates_cache = {} + + for schedule_profile in schedule_profiles: + next_dates = schedule_profile.next_dates + profile_next_dates_cache[schedule_profile.id] = ( + next_dates.get('enable'), + next_dates.get('disable'), + ) + + peer_schedulings_to_update = [] + + for peer_scheduling in peer_scheduling_queryset.iterator(chunk_size=500): + next_enable_at, next_disable_at = profile_next_dates_cache.get( + peer_scheduling.profile_id, + (None, None) + ) + + if not next_enable_at or not next_disable_at: + data['skipped_records'] += 1 + continue + + peer_scheduling.next_scheduled_enable_at = next_enable_at + peer_scheduling.next_scheduled_disable_at = next_disable_at + + peer_schedulings_to_update.append(peer_scheduling) + + if peer_schedulings_to_update: + PeerScheduling.objects.bulk_update( + peer_schedulings_to_update, + ['next_scheduled_enable_at', 'next_scheduled_disable_at'], + batch_size=500 + ) + data['updated_records'] = len(peer_schedulings_to_update) + + return JsonResponse(data) + + def cron_peer_scheduler(request): now = timezone.now() data = { diff --git a/wireguard_webadmin/urls.py b/wireguard_webadmin/urls.py index 4b81286..7ef9d03 100644 --- a/wireguard_webadmin/urls.py +++ b/wireguard_webadmin/urls.py @@ -20,7 +20,7 @@ from django.urls import path from accounts.views import view_create_first_user, view_login, view_logout from api.views import api_instance_info, api_peer_invite, api_peer_list, cron_check_updates, \ cron_update_peer_latest_handshake, peer_info, routerfleet_authenticate_session, routerfleet_get_user_token, \ - wireguard_status, cron_refresh_wireguard_status_cache + wireguard_status, cron_refresh_wireguard_status_cache, cron_calculate_peer_schedules, cron_peer_scheduler from cluster.cluster_api import api_cluster_status, api_get_worker_config_files, api_get_worker_dnsmasq_config, \ api_worker_ping, api_submit_worker_wireguard_stats from cluster.views import cluster_main, cluster_settings, worker_manage @@ -87,6 +87,8 @@ urlpatterns = [ path('api/instance_info/', api_instance_info, name='api_instance_info'), path('api/peer_info/', peer_info, name='api_peer_info'), path('api/peer_invite/', api_peer_invite, name='api_peer_invite'), + path('api/cron_peer_scheduler/', cron_peer_scheduler, name='cron_peer_scheduler'), + path('api/cron_calculate_peer_schedules/', cron_calculate_peer_schedules, name='cron_calculate_peer_schedules'), path('api/cron_refresh_wireguard_status_cache/', cron_refresh_wireguard_status_cache, name='cron_refresh_wireguard_status_cache'), path('api/cron_check_updates/', cron_check_updates, name='cron_check_updates'), path('api/cron_update_peer_latest_handshake/', cron_update_peer_latest_handshake, name='cron_update_peer_latest_handshake'),