#! /usr/bin/env python
# -*- coding: utf-8 -*-
# Программа для просмотра и убийства долгих запросов в Mongo.
 
import logging
import os
import sys
import pymongo
import time
import traceback
import random
import socket
import __main__ as main
import ConfigParser
from pymongo import MongoClient
from optparse import OptionParser

reload(sys)
sys.setdefaultencoding('utf-8')

config = { 'user': "root",
           'password': "",
           'host': "localhost",
           'kill_time': 600,
           'long_time': 0,
           'file': '/etc/mongo-query-kill/main.conf',
           'log': '/var/log/mongo-query-kill/requests.log',
           'graphite_host': 'localhost',
           'graphite_port': 42000 }


def graphiteSender(requests):
    status = False
    sock = None
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(30)
        sock.connect((config['graphite_host'], int(config['graphite_port'])))
        logging.debug('Connect to %s:%s (timeout: %s)' % (config['graphite_host'],
                                                          config['graphite_port'],
                                                          30))
        send_count = 0
        for metric in requests:
            logging.debug('Metric send: %s' % metric)
            sock.send(''.join([metric, '\n']))
            send_count += 1
        status = True
        logging.debug('Sucess send %d params to sender' % send_count)
    except Exception, err:
        logging.warning('Error send data to sender: %s' % err.message)
    finally:
        if sock: sock.close()
        logging.debug('Send status: %s' % status)

def sendToGraphite(longq, killq):
    hostname  = str(socket.getfqdn()).replace('.', '_')
    long_count = 'one_min.{0}.mongodb.long_queries {1} {2}'.format(hostname, len(longq), int(time.time()))
    kill_count = 'one_min.{0}.mongodb.kill_queries {1} {2}'.format(hostname, len(killq), int(time.time()))
    graphite_data = [long_count, kill_count]
    graphiteSender(graphite_data)

def getLongQueries(cnf):
    client = MongoClient('mongodb://{0}:{1}@{2}'.format(cnf['user'], cnf['password'], cnf['host']))
    db = client.admin
    infos = db.current_op()
    longq = [ i for i in infos.get('inprog') if i.get('secs_running') >= int(cnf['long_time']) 
                                               and i.get('active') == True
                                               and i.get('op') in ("query", "getmore") ]
    killq = [ i for i in infos.get('inprog') if i.get('secs_running') >= int(cnf['kill_time'])
                                               and i.get('active') == True
                                               and i.get('op') in ("query", "getmore") ]
    return longq, killq

def doKillQueries(queries, cnf):
    opids = [ q.get('opid') for q in queries ]
    [ logging.critical('killed {0}'.format(i)) for i in opids ]
    client = MongoClient('mongodb://{0}:{1}@{2}'.format(cnf['user'], cnf['password'], cnf['host']))
    [ client['admin']['$cmd.sys.killop'].find_one({'op': opid }) for opid in opids ]
    return

if __name__ == '__main__':
    parser = OptionParser(usage = "usage for help: %prog -h")
    parser.add_option("-u", "--user", dest="user", action="store", default=None,
                      help="пользователь для подключения к mongo")
    parser.add_option("-p", "--password", dest="password", action="store", default=None,
                      help="пароль для подключения к mongodb")
    parser.add_option("-H", "--host", dest="host", action="store", default=None,
                      help="имя сервера с mongo/mongos")
    parser.add_option("-f", "--file", dest="file", action="store", default=None,
                      help="путь до файла с настройками")
    parser.add_option("-t", "--long-time", dest="long_time", action="store", default=None,
                      type="int", help="время, после которого считать процесс долгим(сек)")
    parser.add_option("-k", "--kill-time", dest="kill_time", action="store", default=None,
                      type="int", help="время, после которого убить процесс долгим(сек)")
    parser.add_option("--debug", dest="debug", action="store_true",
                      help="включить debug режим")
    parser.add_option("--kill", dest="kill", action="store_true",
                      help="убивать процессы mongo")
    parser.add_option("--log", dest="log", action="store", default=None,
                      help="имя лога(пишется в неинтерактивном режиме)")        
    parser.add_option("--send-graphite", dest="send_graphite", action="store_true",
                      help="отправлять метрики по запросам в графит")
    parser.add_option("--random-sleep", dest="sleep", action="store_true",
                      help="выполнить скрипт в течении 1 минуты")
  
    opts, args = parser.parse_args()
    log_level = logging.DEBUG if opts.debug else logging.ERROR
    if opts.file: 
        config['file'] = opts.file
    if os.path.exists(config.get('file')): 
        Config = ConfigParser.ConfigParser()
        Config.read(config.get('file'))
        cnf = Config._sections
        config.update(dict(cnf['mongo']))
    else:
        logging.error('Config file {0} not found. Used defaults settings.'.format(config.get('file')))
    parser_opts = vars(opts)
    notnull_opts = [ (i, parser_opts[i]) for i in parser_opts if parser_opts[i] is not None ]
    config.update(dict(notnull_opts))
    # Если программа запущена не в интерактивном режиме, то пишем в лог.
    if sys.stdin.isatty():
        logging.basicConfig(format='%(levelname)s %(asctime)s %(message)s', level=log_level)
    else:
        logging.basicConfig(format='%(levelname)s %(asctime)s %(message)s', level=log_level, filename=config.get('log'), filemode='a')

    if config.get('long_time') > config.get('kill_time'):
        config['kill_time'] = config['long_time']
    logging.debug(config)

    if opts.sleep:
        time.sleep(random.randint(0, 60))

    try:
        longq, killq = getLongQueries(config)
        if config.get('send_graphite', None):
            sendToGraphite(longq, killq)
        if not longq:
            logging.info('Dont found long queries. Exit.')
            os._exit(0)
        logging.critical(time.ctime())
        [ logging.critical('{0}'.format(repr(q).decode("unicode-escape"))) for q in longq ]
        if config.get('kill', None):
            doKillQueries(killq, config)
    except:
        logging.debug(traceback.format_exc().splitlines())
