dbsTest.py

#!/usr/bin/env python

import urllib2,urllib, httplib, sys, re, os, json, phedexSubscription
from xml.dom.minidom import getDOMImplementation
from das_client import get_data
#das_host='https://das.cern.ch'
#das_host='https://cmsweb.cern.ch'
das_host='https://cmsweb-testbed.cern.ch'
#das_host='https://das-dbs3.cern.ch'

def getWorkflowType(url, workflow):
    conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
    r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
    r2=conn.getresponse()
    request = json.loads(r2.read())
    requestType=request['RequestType']
    return requestType


def getRunWhitelist(url, workflow):
    conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
    r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
    r2=conn.getresponse()
    request = json.loads(r2.read())
    runWhitelist=request['RunWhitelist']
    return runWhitelist

def getBlockWhitelist(url, workflow):
    conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
    r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
    r2=conn.getresponse()
    request = json.loads(r2.read())
    BlockWhitelist=request['BlockWhitelist']
    return BlockWhitelist

def getInputDataSet(url, workflow):
    conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
    r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
    r2=conn.getresponse()
    request = json.loads(r2.read())
    while 'exception' in request:
        conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
        r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
        r2=conn.getresponse()
        request = json.loads(r2.read())
    inputDataSets=request['InputDataset']
    if len(inputDataSets)<1:
        print "No InputDataSet for workflow " +workflow
    else:
        return inputDataSets


def duplicateRunLumi(dataset):
    RunlumisChecked={}
    query="file run lumi dataset="+dataset
    das_data = get_data(das_host,query,0,0,0)
    if isinstance(das_data, basestring):
        result = json.loads(das_data)
    else:
        result = das_data
    if result['status'] == 'fail' :
        print 'DAS query failed with reason:',result['reason']
    else:
	preresult=result['data'] 
    for filename in preresult:
     	run=filename['run'][0]['run_number']
	if run not in RunlumisChecked:
		RunlumisChecked[run]=[]
	newLumis=filename['lumi'][0]['number']
	for lumiRange in newLumis:
			newlumiRange=range(lumiRange[0], lumiRange[1]+1)
			for lumi in newlumiRange:
				if lumi in RunlumisChecked[run]:
					return True
				else:
					RunlumisChecked[run].append(lumi)
    return False

def duplicateLumi(dataset):
    lumisChecked=[]
    query="file lumi dataset="+dataset
    das_data = get_data(das_host,query,0,0,0)
    if isinstance(das_data, basestring):
        result = json.loads(das_data)
    else:
        result = das_data
    if result['status'] == 'fail' :
        print 'DAS query failed with reason:',result['reason']
    else:
	preresult=result['data']
	for filename in preresult:
		newLumis=filename['lumi'][0]['number']
		for lumiRange in newLumis:
			newlumiRange=[lumiRange[0]]
			if lumiRange[0]<lumiRange[1]:
				newlumiRange=range(lumiRange[0], lumiRange[1])
			for lumi in newlumiRange:
				if lumi in lumisChecked:
					return True
				else:
					lumisChecked.append(lumi)
	return False

def getRunsInDataset(das_url, dataset):
	query="run dataset="+dataset
  	das_data = get_data(das_url,query,0,0,0)
   	if isinstance(das_data, basestring):
        	result = json.loads(das_data)
	else:
        	result = das_data
	if result['status'] == 'fail' :
        	print 'DAS query failed with reason:',result['reason']
	else:
		runs=[]
		preresult=result['data']
		for run in preresult:
			runs.append(run['run'][0]['run_number'])
		return runs

def getNumberofFilesPerRun(das_url, dataset, run):
	query="file dataset="+dataset+" run="+str(run)+" | count(file.name)"
	das_data = get_data(das_url,query,0,0,0)
   	if isinstance(das_data, basestring):
        	result = json.loads(das_data)
	else:
        	result = das_data
	if result['status'] == 'fail' :
        	print 'DAS query failed with reason:',result['reason']
	else:
		return result['data'][0]['result']['value']

#Return true if there are duplicate evnets , false otherwise
def duplicateEventsMonteCarlo(dataset):
	das_url=das_host
	runs=getRunsInDataset(das_url, dataset)
	for run in runs:
		NumFilesRun=getNumberofFilesPerRun(das_url, dataset, run)
		NumLumis=getRunLumiCountDatasetRun(das_url, dataset, run)
		if NumLumis>NumFilesRun:#It means at least one lumi is split into more than one file
			return True
	return False

#Return the number of events for a given dataset given a runlist
def EventsRunList(das_url, dataset, runlist):
    events=0
    for run in runList:
       events=events+getEventsRun(das_url, dataset, run)
    return events

# Return the number of events in a dataset and in a run
def getEventsRun(das_url, dataset, run):
    query="summary dataset dataset="+dataset+" run="+str(run)+" | grep grep summary.nevents"
    das_data = get_data(das_url,query,0,0,0)
    if isinstance(das_data, basestring):
        result = json.loads(das_data)
    else:
        result = das_data
    if result['status'] == 'fail' :
        print 'DAS query failed with reason:',result['reason']
    else:
	if len(result['data'])==0:#dataset not yet registered in DBS
		return 0
	preresult=result['data'][0]['summary']
	for key in preresult:
		if 'nevents' in key:
			return key['nevents']		
	return -1



#Returns the number of events in a dataset using DAS
def getEventCountDataSet(das_url, dataset):
    query="dataset dataset="+dataset+"  status=* | grep dataset.nevents"
    das_data = get_data(das_url,query,0,0,0)
    if isinstance(das_data, basestring):
        result = json.loads(das_data)
    else:
        result = das_data
    if result['status'] == 'fail' :
        print 'DAS query failed with reason:',result['reason']
    else:
	if len(result['data'])==0:#dataset not yet registered in DBS
		return 0
	preresult=result['data'][0]['dataset']
	for key in preresult:
		if 'nevents' in key:
			return key['nevents']		
	return -1

#Returns a list of runs of a dataset
def getRunsDataset(das_url, dataset):
	runList=[]
	query="run dataset="+dataset+"| grep run.run_number"
	das_data = get_data(das_url,query,0,0,0)
	if isinstance(das_data, basestring):
        	result = json.loads(das_data)
	else:
        	result = das_data
	if result['status'] == 'fail' :
        	print 'DAS query failed with reason:',result['reason']
	else:
		if len(result['data'])==0:#dataset not yet registered in DBS
			return runList
		preresult=result['data']
		for key in preresult:
			if 'run' in key:
				runList.append(key['run'][0]['run_number'])
	return runList

#Returns the number of lumis in a given run for a dataset using DAS.
def getRunLumiCountDatasetRun(das_url, dataset, run):
	lumis=0
	query="summary dataset="+dataset+"  run="+str(run)+" | sum(summary.nlumis)"
	das_data = get_data(das_url,query,0,0,0)
	if isinstance(das_data, basestring):
        	result = json.loads(das_data)
	else:
        	result = das_data
	if result['status'] == 'fail' :
        	print 'DAS query failed with reason:',result['reason']
	else:
		if len(result['data'])==0:#dataset not yet registered in DBS
			return 0
		preresult=result['data']
		for key in preresult:
		    if 'result' in key:
			return key['result']['value']		
		return -1

# Get the number of unique lumis in a dataset
def getRunLumiCountDataset(das_url, dataset):
        lumis=0
        query="summary dataset="+dataset+" | grep summary.nlumis"
        das_data = get_data(das_url,query,0,0,0)
        if isinstance(das_data, basestring):
                result = json.loads(das_data)
        else:
                result = das_data
        if result['status'] == 'fail' :
                print 'DAS query failed with reason:',result['reason']
        else:
                if len(result['data'])==0:#dataset not yet registered in DBS
                        return 0
                preresult=result['data'][0]
                for key in preresult:
                    if 'summary' in key:
                        return preresult['summary'][0]['nlumis']               
                return -1


#Returns the number of lumis in a dataset and in a given runlist
def getRunLumiCountDatasetListDAS(das_url,dataset, runlist):
    lumis=0
    runChunks=chunks(runlist,30)
    for runList in runChunks:
	lumis=lumis+getRunLumiCountDatasetListLimited(das_url,dataset, runList)
    return lumis

#Returns the number of lumis in a dataset and in a given runlist but with a limit of 30 runs	
def getRunLumiCountDatasetListLimited(das_url,dataset, runlist):
    lumis=0
    query="summary dataset="+dataset+" run in "+str(runlist)+ " | grep summary.nlumis"
    das_data = get_data(das_url,query,0,0,0)
    if isinstance(das_data, basestring):
       	result = json.loads(das_data)
    else:
       	result = das_data
    if result['status'] == 'fail' :
        print 'DAS query failed with reason:',result['reason']
    else:
	if len(result['data'])==0:#dataset not yet registered in DBS
		return 0
	preresult=result['data'][0]
	for key in preresult:
		if 'summary' in key:
			return preresult['summary'][0]['nlumis']		
	return -1
	

#GEN check generator true if everything is fine, false otherwise
def checkCorrectLumisEventGEN(dataset):
	das_url=das_host
	numlumis=getRunLumiCountDataset(das_url, dataset)
	numEvents=getEventCountDataSet(das_url, dataset)
	if numlumis>=numEvents/300.0:
		return True
	else:
		return False
  

# SPlits a list of chunks of size(n)
def chunks(lis, n):
    return [lis[i:i+n] for i in range(0, len(lis), n)]
# Return the number of events in a block using DAS
def getEventsBlock(das_url, block_name):
    query="block="+block_name+"  | grep block.nevents"
    das_data = get_data(das_url,query,0,0,0)
    if isinstance(das_data, basestring):
        result = json.loads(das_data)
    else:
        result = das_data
    if result['status'] == 'fail' :
        print 'DAS query failed with reason:',result['reason']
    else:
	if len(result['data'])==0:#dataset not yet registered in DBS
		return 0
	preresult=result['data'][0]['block']
	for key in preresult:
		if 'nevents' in key:
			return key['nevents']		
	return -1
    return 0

#Return the number of events for a given dataset given a blocklist
def EventsBlockList(das_url, dataset, blocklist):
    events=0
    if len(blocklist)==0:
        return getEventCountDataSet(das_url, dataset)
    for block in blocklist:
	events=events+getEventsBlock(das_url, block)
    return events


def getRunLumiCountDatasetBlockList(dataset,blockList):
    lumis=0
    querry="./dbssql --limit=1000000 --input='find block, run, lumi where dataset="+dataset+"' |egrep '("
    for block in blockList:
        querry=querry+block+ "|"
    querry=querry+blockList[0]+")'"
    querry=querry+"| awk '{print $2, $3}' | sort | uniq | wc -l"
    output=os.popen(querry).read()
    if not output:
        lumis=lumis
    try:
        lumis=lumis+int(output)
    except ValueError:
        lumis=lumis
    return lumis

def handleTaskChain(request):
    # Check if it's MC from scratch
    if 'RequestNumEvents' in request['Task1']:
        if request['Task1']['RequestNumEvents'] is not None:
            return request['Task1']['RequestNumEvents']

    blockWhitelist = blockBlacklist = runWhitelist = runBlacklist = []
    if 'InputDataset' in request['Task1']:
        inputDataSet=request['Task1']['InputDataset']
        if 'BlockWhitelist' in request['Task1']:
            blockWhitelist=request['Task1']['BlockWhitelist']
        if 'BlockBlacklist' in request['Task1']:
            blockBlacklist=request['Task1']['BlockBlacklist']
        if 'RunWhitelist' in request['Task1']:
            runWhitelist=request['Task1']['RunWhitelist']
        if 'RunBlacklist' in request['Task1']:
            runBlacklist=request['Task1']['RunBlacklist']

        if len(blockWhitelist)>0:
            return getRunLumiCountDatasetBlockList(das_host, inputDataSet,blockWhitelist)
        if len(blockBlacklist)>0:
            return getRunLumiCountDataset(das_host,inputDataSet)-getRunLumiCountDatasetBlockList(das_host, inputDataSet,blockBlacklist)
        if len(runWhitelist)>0:
            return getRunLumiCountDatasetListDAS(das_host, inputDataSet, runWhitelist)
        else:
            return getRunLumiCountDataset(das_host,inputDataSet)

### TODO: implement multi white/black list
#        if len(blockWhitelist)>0 and len(runWhitelist)>0:
#            print "Hey, you have block and run white list :-D"
#            return getRunLumiCountDatasetBlockList(inputDataSet,BlockWhitelist)
#        elif len(blockWhitelist)>0 and len(runWhitelist)==0:
#            print "Hey, you have block white list but NOT run white list :-D"
#        elif len(blockWhitelist)==0 and len(runWhitelist)>0:
#            print "Hey, you have NO block white list but you do have run white list :-D"
#            return getRunLumiCountDatasetList(inputDataSet, runWhitelist)
#        elif len(blockWhitelist)==0 and len(runWhitelist)==0:
#            print "Hey, you have NO block and run white list :-D"
#
#        if len(BlockBlacklist)>0 and len(runBlacklist)>0:
#            print "Hey, you have block and run black list :-("
#            return getRunLumiCountDataset(inputDataSet)-getRunLumiCountDatasetBlockList(inputDataSet,BlockBlacklist)
#        elif len(BlockBlacklist)>0 and len(runBlacklist)==0:
#            print "Hey, you have block black list but NOT run black list :-("
#        elif len(BlockBlacklist)==0 and len(runBlacklist)>0:
#            print "Hey, you have NO block black list but you do have run black list :-("
#        elif len(BlockBlacklist)==0 and len(runBlacklist)==0:
#            print "Hey, you have NO block and run black list :-("

def getInputEvents(url, workflow):
    conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
    r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
    r2=conn.getresponse()
    request = json.loads(r2.read())
    while 'exception' in request:
        conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
        r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
        r2=conn.getresponse()
        request = json.loads(r2.read())
    requestType=request['RequestType']
    if requestType=='MonteCarlo' or requestType=='LHEStepZero':
        if 'RequestNumEvents' in request:
            if request['RequestNumEvents']>0:
                return request['RequestNumEvents']
        if 'RequestSizeEvents' in request:
            return request['RequestSizeEvents']
        else:
            return 0
    if requestType == 'TaskChain':
        return handleTaskChain(request)

    BlockWhitelist=request['BlockWhitelist']
    BlockBlacklist=request['BlockBlacklist']
    inputDataSet=request['InputDataset']
    runWhitelist=[]
    if 'RunWhitelist' in request:
	runWhitelist=request['RunWhitelist']
    if requestType=='ReReco':
        if len(BlockWhitelist)>0:
            return getRunLumiCountDatasetBlockList(inputDataSet,BlockWhitelist)
        if len(BlockBlacklist)>0:
            return getRunLumiCountDataset(inputDataSet)-getRunLumiCountDatasetBlockList(inputDataSet,BlockBlacklist)
        if len(runWhitelist)>0:
            return getRunLumiCountDatasetListDAS(das_host, inputDataSet, runWhitelist)
        else:
            return getRunLumiCountDataset(inputDataSet)
    events=getEventCountDataSet(das_host, inputDataSet)
    if len(BlockBlacklist)>0:
        events=events-EventsBlockList(inputDataSet, BlockBlacklist)
    if len(runWhitelist)>0:
        events=EventsRunList(das_host, inputDataSet, runWhitelist)
    if len(BlockWhitelist)>0:
        events=EventsBlockList(das_host, inputDataSet, BlockWhitelist)
    if 'FilterEfficiency' in request.keys():
        return float(request['FilterEfficiency'])*events
    else:
        return events

def getOutputEvents(url, workflow, dataset):
    conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
    r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
    r2=conn.getresponse()
    request = json.loads(r2.read())
    while 'exception' in request:
        conn  =  httplib.HTTPSConnection(url, cert_file = os.getenv('X509_USER_PROXY'), key_file = os.getenv('X509_USER_PROXY'))
        r1=conn.request("GET",'/reqmgr/reqMgr/request?requestName='+workflow)
        r2=conn.getresponse()
        request = json.loads(r2.read())
    requestType=request['RequestType']
    if requestType=='ReReco':
        return getRunLumiCountDataset(das_host, dataset)
    elif requestType=='TaskChain':
        if 'InputDataset' in request['Task1']:
             return getRunLumiCountDataset(das_host, dataset)
        else:
             return getEventCountDataSet(das_host, dataset)
    else:
        return getEventCountDataSet(das_host, dataset)

  
def main():
    args=sys.argv[1:]
    if not len(args)==1:
        print "usage:dbsTest workflowname"
        sys.exit(0)
    workflow=args[0]
    url='cmsweb.cern.ch'
    outputDataSets=phedexSubscription.outputdatasetsWorkflow(url, workflow)
    inputEvents=getInputEvents(url, workflow)
    for dataset in outputDataSets:
        outputEvents=getOutputEvents(url, workflow, dataset)
        if inputEvents!=0:
            print dataset+" match: "+str(outputEvents/float(inputEvents)*100) +"%"
        else:
            print "Input Events 0"
    sys.exit(0);

if __name__ == "__main__":
    main()