#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2009-2011 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 Condor machine slots and submitter 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(hours=1))
DEFAULT_END_DT = str(datetime.utcnow()-UTC_DIFF)

verbose = False

def print_user(user,start,end):
	for item in db['samples.submitter'].find({"sn":{'$regex':'^'+user}, 'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}):
		print item['sn'],"\t",item['ts']-UTC_DIFF, "\t",item['jr'],"\t",item['jh'],"\t",item['ji']

def print_resource(resource,start,end):
	for item in db['samples.machine'].find({"mn":{'$regex':resource}, 'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}):
		print item['mn'],"\t",item['ts']-UTC_DIFF, "\t","%s/%s" % (item['ar'],item['os']),"\t",item['ki'],"\t",str(item['la'])[:5],"\t",item['st']

def print_users():
		for user in db['samples.submitter'].distinct('sn'):
			print user

def print_usergroups():
	# TODO: revisit with mongodb group()
	mnlist = db['samples.submitter'].distinct('mn')
	for mn in mnlist:
		snlist = db['samples.submitter'].find({'mn': mn}).distinct('sn')
		for sn in snlist:
			print '%s/%s' % (sn,mn)

def print_resourcelist():
	for item in db['samples.machine'].distinct('mn'):
		print item

parser = OptionParser(description='Query Condor ODS for statistics')
parser.add_option('-v','--verbose', action="store_true",default=False, help='enable logging')
parser.add_option('-s','--server', action="store", nargs='?', dest='server',
                    default='localhost',
                    help='mongodb database server location: e.g., somehost, localhost:2011')
parser.add_option('--u','--user', dest="user", help='stats for a single submitter: user,timestamp,running,held,idle')
parser.add_option('--r','--resource', dest="resource", help='stats for a single resource: slot,timestamp,keyboard idle,load average,status')
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('--ul','--userlist', action="store_true",dest="userlist", default=False, help='list all submitters')
parser.add_option('--ugl','--usergrouplist', action="store_true",dest='usergrouplist',default=False, help='list all submitter groups')
parser.add_option('--rl','--resourcelist', action="store_true",dest='resourcelist',default=False, help='list all resources')
(options, args) =  parser.parse_args()

verbose = options.verbose

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

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

if options.user:
	print_user(options.user,options.start,options.end)
elif options.resource:
	print_resource(options.resource,options.start,options.end)
elif options.userlist:
	print_users()
elif options.usergrouplist:
	print_usergroups()
elif options.resourcelist:
	print_resourcelist()
else:
	parser.print_help()
