Skip to content
Snippets Groups Projects
rc_util.py 1.14 KiB
Newer Older
Bo-Chun Chen's avatar
Bo-Chun Chen committed
from rc_rmq import RCRMQ
import json

rc_rmq = RCRMQ({'exchange': 'RegUsr', 'exchange_type': 'topic'})
tasks = {'ohpc_account': None, 'ood_account': None, 'slurm_account': None}
Bo-Chun Chen's avatar
Bo-Chun Chen committed

def add_account(username, full='', reason=''):
  rc_rmq.publish_msg({
    'routing_key': 'request.' + username,
Bo-Chun Chen's avatar
Bo-Chun Chen committed
    'msg': {
      "username": username,
      "fullname": full,
      "reason": reason
    }
  })
  rc_rmq.disconnect()
Bo-Chun Chen's avatar
Bo-Chun Chen committed

def worker(ch, method, properties, body):
    msg = json.loads(body)
    task = msg['task']
    tasks[task] = msg['success']
    print("Got msg: {}({})".format(msg['task'], msg['success']))
Bo-Chun Chen's avatar
Bo-Chun Chen committed

    # Check if all tasks are done
    done = True
    for key, status in tasks.items():
        if not status:
Bo-Chun Chen's avatar
Bo-Chun Chen committed
            print("{} is not done yet.".format(key))
Bo-Chun Chen's avatar
Bo-Chun Chen committed
            done = False
Bo-Chun Chen's avatar
Bo-Chun Chen committed
    if done:
Bo-Chun Chen's avatar
Bo-Chun Chen committed
        rc_rmq.stop_consume()
        rc_rmq.delete_queue()
def consume(username, callback=worker, debug=False):
Bo-Chun Chen's avatar
Bo-Chun Chen committed
    if debug:
        sleep(5)
    else:
Bo-Chun Chen's avatar
Bo-Chun Chen committed
        rc_rmq.start_consume({
Bo-Chun Chen's avatar
Bo-Chun Chen committed
            'queue': username,
            'routing_key': 'confirm.' + username,
            'cb': callback
Bo-Chun Chen's avatar
Bo-Chun Chen committed
        })
        rc_rmq.disconnect()
Bo-Chun Chen's avatar
Bo-Chun Chen committed
    return { 'success' : True }