neptune/core/management/commands/scheduling.py

53 lines
1.6 KiB
Python

import asyncio
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from asgiref.sync import sync_to_async
from django.core.management.base import BaseCommand
from core.db.storage import db
from core.lib.parsing import QueryError
from core.lib.rules import NotificationRuleData
from core.models import NotificationRule
from core.util import logs
log = logs.get_logger("scheduling")
INTERVALS = [5, 60, 900, 1800, 3600, 14400, 86400]
async def job(interval_seconds):
"""
Run all schedules matching the given interval.
:param interval_seconds: The interval to run.
"""
matching_rules = await sync_to_async(list)(
NotificationRule.objects.filter(enabled=True, interval=interval_seconds)
)
for rule in matching_rules:
log.debug(f"Running rule {rule}")
try:
rule = NotificationRuleData(rule.user, rule, db=db)
await rule.run_schedule()
# results = await db.schedule_query_results(rule.user, rule)
except QueryError as e:
log.error(f"Error running rule {rule}: {e}")
class Command(BaseCommand):
def handle(self, *args, **options):
"""
Start the scheduling process.
"""
scheduler = AsyncIOScheduler()
for interval in INTERVALS:
log.debug(f"Scheduling {interval} second job")
scheduler.add_job(job, "interval", seconds=interval, args=[interval])
scheduler.start()
loop = asyncio.get_event_loop()
try:
loop.run_forever()
except (KeyboardInterrupt, SystemExit):
log.info("Process terminating")
finally:
loop.close()