#!/usr/bin/env python
# Copyright European Organization for Nuclear Research (CERN) 2014
#
# 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:
# - Vincent Garonne, <vincent.garonne@cern.ch>, 2014

'''
Probe to check replicas in transient states.
'''

import sys

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

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

if __name__ == "__main__":
    try:
        session = get_session()
        result = session.execute('select  /*+ index_ffs(replicas REPLICAS_TOMBSTONE_IDX) */  count(1) from atlas_rucio.replicas where tombstone is not null').fetchone()[0]
        monitor.record_gauge(stat='replicas.unlocked.total', value=result)
        print result

        query = '''select  /*+ index_ffs(replicas REPLICAS_TOMBSTONE_IDX) */  count(1), state, atlas_rucio.id2rse(rse_id)
                   from atlas_rucio.replicas
                   where tombstone is not null and tombstone < sys_extract_utc(systimestamp)
                   group by state, rse_id
                '''
        results = session.execute(query)
        for count, state, rse in results:
            print rse, state, count
            monitor.record_gauge(stat='replicas.unlocked.expired.%(rse)s.%(state)s' % locals(),   value=count)

        query = '''select  /*+ index(replicas, REPLICAS_STATE_IDX) */  count(*), state, atlas_rucio.id2rse(rse_id)
                   from atlas_rucio.replicas r where (case when STATE != 'A' then RSE_ID END) is not null
                   group by state, rse_id'''
        results = session.execute(query)
        for count, state, rse in results:
            print rse, state, count
            monitor.record_gauge(stat='replicas.transient_state.%(state)s.rse.%(rse)s' % locals(),   value=count)

    except:
        sys.exit(UNKNOWN)
    sys.exit(OK)
