#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2009-2012 Red Hat, Inc.
#
# 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# This utility script is designed to query the mongodb database
# for sampled records of scheduler job totals.

# uses pymongo - http://pypi.python.org/pypi/pymongo/
import pymongo
from datetime import timedelta, datetime
from sys import exit, argv
import time, pwd, os
from optparse import OptionParser
from dateutil import *
from dateutil.parser import *

# NOTE: the 1.9 pymongo driver does implicit localtime conversion when
# dates are used in a query so we must compensate, eventhough the dates
# ARE already stored in localtime in 1.6.4 mongodb
# mongodb 1.7+ uses ISODate to address this
UTC_DIFF = datetime.utcnow() - datetime.now()
DEFAULT_START_DT = str(datetime.utcnow()-UTC_DIFF-timedelta(minutes=15))
DEFAULT_END_DT = str(datetime.utcnow()-UTC_DIFF)

verbose = False

def print_schedulers():
    print "SCHEDULER".ljust(40),"RUNNING".ljust(8),"IDLE".ljust(8),"HELD".ljust(8),"REMOVED".ljust(8)
    for schedd in raw['ads'].find({'MyType': 'Scheduler'}):
        print schedd['Name'].ljust(40),str(schedd['TotalRunningJobs']).ljust(8),str(schedd['TotalIdleJobs']).ljust(8),\
            str(schedd['TotalHeldJobs']).ljust(8),str(schedd['TotalRemovedJobs']).ljust(8)

def print_stats(start,end,scheduler):
    print "TIMESTAMP".ljust(30),"RUNNING".ljust(8),"IDLE".ljust(8),"HELD".ljust(8),"REMOVED".ljust(8)
    for schedd in stats['samples.scheduler'].find({'n':scheduler,'ts':{'$gte': parse(start)+UTC_DIFF, \
            '$lt': parse(end)+UTC_DIFF}}).sort('ts',pymongo.DESCENDING):
        print str(schedd['ts']-UTC_DIFF).ljust(30), str(schedd['trun']).ljust(8),str(schedd['tij']).ljust(8),\
            str(schedd['thj']).ljust(8),str(schedd['trem']).ljust(8)

parser = OptionParser(description='Query Condor ODS for scheduler job totals')
parser.add_option('-v','--verbose', action="store_true",default=False, help='enable logging')
parser.add_option('-s','--server', action="store", dest='server',
                    default='localhost',
                    help='mongodb database server location: e.g., somehost, localhost:2011')
parser.add_option('-f','--from', dest="start", help='records from datetime in ISO8601 format e.g., \'2011-09-29 12:03\'', default=DEFAULT_START_DT)
parser.add_option('-t','--to', dest="end", help='records to datetime in ISO8601 format e.g., \'2011-09-30T17:16\'',default=DEFAULT_END_DT)
parser.add_option('-S','--schedd', dest="scheduler", help='report using specific scheduler in start to end range')
(options, args) =  parser.parse_args()

verbose = options.verbose

try:
    connection = pymongo.Connection(options.server)
    stats = connection.condor_stats
    raw = connection.condor_raw
except Exception, e:
    print e
    exit(1)

if verbose:
    print 'from:\t', options.start
    print 'to:\t', options.end

if options.scheduler:
    print_stats(options.start,options.end,options.scheduler)
else:
    print_schedulers()
