getWQStatusByWorkflow.py

"""
__getWQStatusByWorkflow.py__

Given a workflow name, it collects the status of the workqueue elements status in:
 a) global workqueue
 b) local workqueue
 c) local workqueue_inbox

Created on September 23, 2016.
@author: amaltaro
"""

import sys
import os
from pprint import pprint

from WMCore.Configuration import loadConfigurationFile
from WMCore.WorkQueue.WorkQueueBackend import WorkQueueBackend


def createElementsSummary(reqName, elements, queueUrl):
    """
    Print the local couchdb situation based on the WQE status
    """
    print("\nSummary for %s and request %s" % (queueUrl, reqName))
    summary = {'numberOfElements': len(elements)}
    for elem in elements:
        summary.setdefault(elem['Status'], {})
        if elem['ChildQueueUrl'] not in summary[elem['Status']]:
            summary[elem['Status']][elem['ChildQueueUrl']] = 0
        summary[elem['Status']][elem['ChildQueueUrl']] += 1
    pprint(summary)


def main():
    """
    Whatever
    """
    if 'WMAGENT_CONFIG' not in os.environ:
        os.environ['WMAGENT_CONFIG'] = '/data/srv/wmagent/current/config/wmagent/config.py'
    config = loadConfigurationFile(os.environ["WMAGENT_CONFIG"])

    if len(sys.argv) != 2:
        print("You must provide a request name")
        sys.exit(1)

    reqName = sys.argv[1]

    globalWQBackend = WorkQueueBackend(config.WorkloadSummary.couchurl, db_name="workqueue")
    localWQBackend = WorkQueueBackend(config.WorkQueueManager.couchurl, db_name="workqueue")
    localWQInbox = WorkQueueBackend(config.WorkQueueManager.couchurl, db_name="workqueue_inbox")

    gqDocIDs = globalWQBackend.getElements(RequestName=reqName)
    localDocIDs = localWQBackend.getElements(RequestName=reqName)
    localInboxDocIDs = localWQInbox.getElements(RequestName=reqName)

    createElementsSummary(reqName, gqDocIDs, globalWQBackend.queueUrl)
    createElementsSummary(reqName, localDocIDs, localWQBackend.queueUrl)
    createElementsSummary(reqName, localInboxDocIDs, localWQInbox.queueUrl)
    
    sys.exit(0)

if __name__ == "__main__":
    sys.exit(main())