django-orchestra/orchestra/contrib/tasks/bin/orchestra-beat

85 lines
2.9 KiB
Python
Executable File

#!/usr/bin/env python3
# High performance alternative to beat management command
#
# USAGE: beat /path/to/project/manage.py
import json
import os
import re
import sys
from datetime import datetime
from orchestra.utils import db
from orchestra.utils.python import import_class
from orchestra.utils.sys import run, join
from celery.schedules import crontab_parser as CrontabParser
def get_settings_file(manage):
with open(manage, 'r') as handler:
regex = re.compile(r'"DJANGO_SETTINGS_MODULE"\s*,\s*"([^"]+)"')
for line in handler.readlines():
match = regex.search(line)
if match:
settings_module = match.groups()[0]
settings_file = os.path.join(*settings_module.split('.')) + '.py'
settings_file = os.path.join(os.path.dirname(manage), settings_file)
return settings_file
raise ValueError("settings module not found in %s" % manage)
def get_tasks(manage):
settings_file = get_settings_file(manage)
settings = db.get_settings(settings_file)
try:
conn = db.get_connection(settings)
except:
sys.stdout.write("ERROR")
sys.stderr.write("I am unable to connect to the database\n")
sys.exit(1)
script, settings_file = sys.argv[:2]
enabled = 1 if 'sqlite' in settings['ENGINE'] else True
query = (
"SELECT c.minute, c.hour, c.day_of_week, c.day_of_month, c.month_of_year, p.id "
"FROM djcelery_periodictask as p, djcelery_crontabschedule as c "
"WHERE p.crontab_id = c.id AND p.enabled = {}"
).format(enabled)
tasks = db.run_query(conn, query)
conn.close()
return tasks
def is_due(now, minute, hour, day_of_week, day_of_month, month_of_year):
n_minute, n_hour, n_day_of_week, n_day_of_month, n_month_of_year = now
return (
n_minute in CrontabParser(60).parse(minute) and
n_hour in CrontabParser(24).parse(hour) and
n_day_of_week in CrontabParser(7).parse(day_of_week) and
n_day_of_month in CrontabParser(31, 1).parse(day_of_month) and
n_month_of_year in CrontabParser(12, 1).parse(month_of_year)
)
if __name__ == "__main__":
manage = sys.argv[1]
now = datetime.utcnow()
now = tuple(map(int, now.strftime("%M %H %w %d %m").split()))
procs = []
for minute, hour, day_of_week, day_of_month, month_of_year, task_id in get_tasks(manage):
if is_due(now, minute, hour, day_of_week, day_of_month, month_of_year):
command = 'python3 -W ignore::DeprecationWarning {manage} runtask {task_id}'.format(
manage=manage, task_id=task_id)
proc = run(command, async=True)
procs.append(proc)
code = 0
for proc in procs:
result = join(proc)
sys.stdout.write(result.stdout.decode('utf8'))
sys.stderr.write(result.stderr.decode('utf8'))
if result.return_code != 0:
code = result.return_code
sys.exit(code)