#!/usr/bin/env python
# Copyright European Organization for Nuclear Research (CERN) 2013
#
# Licensed under the Apache License, Version 2.0 (the "License");
# You may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Authors:
# - Mario Lassnig, <mario.lassnig@cern.ch>, 2013

'''
Probe to check the queues of the transfer service
'''

import sys

from rucio.core import monitor
from rucio.db.session import get_session

# Exit statuses
UNKNOWN, OK, WARNING, CRITICAL = -1, 0, 1, 2

queues = ["""SELECT
CASE
 WHEN state = 'Q' THEN 'requests.history.queued'
 WHEN state = 'S' THEN 'requests.history.submitted'
 WHEN state = 'F' THEN 'requests.history.failed'
 WHEN state = 'D' THEN 'requests.history.done'
 WHEN state = 'L' THEN 'requests.history.lost'
 ELSE state
END state_desc ,
num_rows
FROM
(
select  /*+ INDEX(REQUESTS_HISTORY REQUESTS_H_TYP_STA_IDX) */ state, count(*) num_rows
FROM atlas_rucio.requests_history
WHERE request_type='T'
GROUP BY state
)""", """SELECT
CASE
 WHEN state = 'Q' THEN 'requests.queued'
 WHEN state = 'S' THEN 'requests.submitted'
 WHEN state = 'F' THEN 'requests.failed'
 WHEN state = 'D' THEN 'requests.done'
 WHEN state = 'L' THEN 'requests.lost'
 ELSE state
END state_desc ,
num_rows
FROM
(
select state, count(*) num_rows
FROM atlas_rucio.requests
WHERE request_type='T'
GROUP BY state
)"""]

if __name__ == "__main__":
    try:
        session = get_session()
        for q in queues:
            result = session.execute(q).fetchall()
            for r in result:
                print 'queues.%s' % r[0], r[1],
                monitor.record_gauge(stat='queues.%s' % r[0], value=r[1])
    except:
        sys.exit(UNKNOWN)
    sys.exit(OK)
