#!/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-2015
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2014
# - Wen Guan, <wen.guan@cern.ch>, 2015

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

import sys

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

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

active_queue = """SELECT
CASE
 WHEN state = 'S' THEN 'queues.requests.submitted.' || replace(activity, ' ', '_') || '.' || nvl(replace(substr(external_host, 9, length(external_host)-13), '.', '_'), 'no_fts_host')
 WHEN state = 'Q' THEN 'queues.requests.queued.' || replace(activity, ' ', '_') || '.' || nvl(replace(substr(external_host, 9, length(external_host)-13), '.', '_'), 'no_fts_host')
 WHEN state = 'F' THEN 'queues.requests.failed.' || replace(activity, ' ', '_') || '.' || nvl(replace(substr(external_host, 9, length(external_host)-13), '.', '_'), 'no_fts_host')
 WHEN state = 'D' THEN 'queues.requests.done.' || replace(activity, ' ', '_') || '.' || nvl(replace(substr(external_host, 9, length(external_host)-13), '.', '_'), 'no_fts_host')
 WHEN state = 'L' THEN 'queues.requests.lost.' || replace(activity, ' ', '_') || '.' || nvl(replace(substr(external_host, 9, length(external_host)-13), '.', '_'), 'no_fts_host')
 WHEN state = 'W' THEN 'queues.requests.waiting.' || replace(activity, ' ', '_') || '.' || nvl(replace(substr(external_host, 9, length(external_host)-13), '.', '_'), 'null_fts_host')
 ELSE state
END state_desc,
num_rows
FROM
(
select state, count(*) num_rows, activity, external_host
FROM atlas_rucio.requests
GROUP BY state, activity, external_host
)"""

if __name__ == "__main__":
    try:
        session = get_session()
        for k in session.execute(active_queue).fetchall():
            print k[0], k[1],
            monitor.record_gauge(stat=k[0], value=k[1])
    except:
        sys.exit(UNKNOWN)
    sys.exit(OK)
