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

#!/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)