dramatiq

#!/usr/bin/env python
import os
import importlib

from mist.api import config

BROKER = 'mist.api.dramatiq_app'
QUEUES = os.getenv('QUEUES', '').split(',') or [
    'default',
    'dramatiq_provisioning',
    'dramatiq_scripts',
    'dramatiq_schedules',
    'dramatiq_rules',
    'dramatiq_sessions',
    'dramatiq_polling',
    'dramatiq_mappings',
    'dramatiq_ping_probe',
    'dramatiq_ssh_probe',
    'dramatiq_machines',
    'dramatiq_clusters',
    'dramatiq_networks',
    'dramatiq_zones',
    'dramatiq_volumes',
    'dramatiq_secrets',
    'dramatiq_buckets',
    'dramatiq_tags']

POSSIBLE_MODULES = []  # 'mist.api.dramatiq_tasks', 'mist.rbac.dramatiq_tasks']

command = 'dramatiq'
processes = os.getenv('DRAMATIQ_PROCESSES')
if processes:
    command += f' --processes {processes}'
threads = os.getenv('DRAMATIQ_THREADS')
if threads:
    command += f' --threads {threads}'

for root, dirs, files in os.walk("src"):
    for file in files:
        if file.endswith('tasks.py'):
            module = os.path.join(
                root, file[:-3]).split('src/')[1].replace('/', '.')
            POSSIBLE_MODULES.append(module)
            print(module)

for root, dirs, files in os.walk("/opt"):
    for file in files:
        if file.endswith('tasks.py'):
            print('.'.join(os.path.join(root, file[:-3]).split('/')[3:]))


modules = []

for module in POSSIBLE_MODULES:
    try:
        importlib.import_module(module)
    except ModuleNotFoundError:
        continue
    modules.append(module)
    print(f"  * {module}")

for plugin in config.PLUGINS:
    module = 'mist.' + plugin + '.tasks'
    try:
        importlib.import_module(module)
    except ModuleNotFoundError:
        continue
    modules.append(module)
    print(f"  * {module}")

print(
    f"Will execute command: {command} {BROKER} {' '.join(modules)} -Q {' '.join(QUEUES)}")

os.system(f"{command} {BROKER} {' '.join(modules)} -Q {' '.join(QUEUES)}")