Python CMSCouch.Database类代码示例

def requestDetails(requestName):
    """ Adds details from the Couch document as well as the database """
    WMCore.Lexicon.identifier(requestName)
    request = GetRequest.getRequestDetails(requestName)
    helper = loadWorkload(request)
    schema = helper.data.request.schema.dictionary_whole_tree_()
    # take the stuff from the DB preferentially
    schema.update(request)
    task = helper.getTopLevelTask()[0]
    
    schema['Site Whitelist']  = task.siteWhitelist()
    schema['Site Blacklist']  = task.siteBlacklist()
    schema['MergedLFNBase']   = str(helper.getMergedLFNBase())
    schema['UnmergedLFNBase'] = str(helper.getUnmergedLFNBase())
    schema['Campaign']        = str(helper.getCampaign()) 
    schema['AcquisitionEra']  = str(helper.getAcquisitionEra())
    if schema['SoftwareVersions'] == ['DEPRECATED']:
        schema['SoftwareVersions'] = helper.getCMSSWVersions()

    # Check in the CouchWorkloadDBName if not present
    schema.setdefault("CouchWorkloadDBName", "reqmgr_workload_cache")

    # get DbsUrl from CouchDB
    if schema.get("CouchWorkloadDBName", None) and schema.get("CouchURL", None):
        couchDb = Database(schema["CouchWorkloadDBName"], schema["CouchURL"])
        couchReq = couchDb.document(requestName)
        schema["DbsUrl"] = couchReq.get("DbsUrl", None)
        
    # https://github.com/dmwm/WMCore/issues/4588
    schema["SubscriptionInformation"] = helper.getSubscriptionInformation()
    return schema

def main():
    config = loadConfigurationFile(os.environ['WMAGENT_CONFIG'])
    config.CoreDatabase.dialect = 'oracle'
    init = WMInit()
    init.setDatabaseConnection(config.CoreDatabase.connectUrl,
                               config.CoreDatabase.dialect)
    couchDB = Database('wmagent_jobdump/fwjrs', '')
    couchDB2 = Database('wmagent_jobdump/jobs', '')
    myThread = threading.currentThread()
    daofactory = DAOFactory(package = "WMCore.WMBS",
                            logger = logging,
                            dbinterface = myThread.dbi)
    getJobsDAO = daofactory(classname = "Jobs.GetAllJobs")
    completedJobs = getJobsDAO.execute(state = 'complete')
    candidates = []
    while len(completedJobs):
        candidates = []
        chunk = completedJobs[:500]
        completedJobs = completedJobs[500:]
        result = couchDB.loadView('FWJRDump', 'outputByJobID', keys = chunk)
        rows = result['rows']
        for entry in rows:
            candidates.append(entry['key'])
        for jobId in candidates:
            doc = couchDB2.document(str(jobId))
            last = max(map(int, doc['states'].keys()))
            lastState = doc['states'][str(last)]['newstate']
            if lastState == 'success':
                print jobId

def getFileInformation(workflow, lfn, outModule):

    # Connect to the FWJR DB
    fwjrDB = Database('wmagent_jobdump/fwjrs', 'http://dummy.cern.ch:5984')

    result = fwjrDB.loadView('FWJRDump', 'jobsByOutputLFN', {'include_docs' : True}, [[workflow, lfn]])
    if result['rows']:
        fwjrDoc = result['rows'][0]['doc']
        fwjrInfo = fwjrDoc['fwjr']
        for step in fwjrInfo['steps']:
            if step == 'cmsRun1':
                if outModule not in fwjrInfo['steps'][step]['output']:
                    print "WARNING: No output module %s in this job" % outModule
                    return
                outModuleInfo = fwjrInfo['steps'][step]['output'][outModule]
                for fileInfo in outModuleInfo:
                    if fileInfo['lfn'] == lfn:
                        print "File information, %s" % fileInfo['lfn']
                        print "Run/Lumis:"
                        for run in fileInfo['runs']:
                            print 'Run: %s, Lumi range: %s-%s' % (run, fileInfo['runs'][run][0], fileInfo['runs'][run][1])
                        print "Number of Events: %s" % fileInfo['events']
                        print "Filesize (bytes): %.1f" % (float(fileInfo['size']))
                        print "Adler32 Checksum: %s" % fileInfo['checksums']['adler32']
    else:
        print "WARNING: No file info in CouchDB"

    return

def checkWorkQueue(requestName):
    result = {'ActiveAgents' : {},
              'ElementsRunning' : 0,
              'ElementsAcquired' : 0,
              'ElementsAvailable' : 0,
              'ElementsDone' : 0}
    x = Database('workqueue', 'https://cmsweb.cern.ch/couchdb')
    y = x.loadView('WorkQueue', 'elementsByParent', {'include_docs' : True}, [requestName])
    for entry in y['rows']:
        doc = entry['doc']
        element = doc['WMCore.WorkQueue.DataStructs.WorkQueueElement.WorkQueueElement']
        status = element['Status']
        if status == 'Running':
            result['ElementsRunning'] += 1
        elif status == 'Acquired':
            result['ElementsAcquired'] += 1
        elif status == 'Available':
            result['ElementsAvailable'] += 1
        elif status == 'Done':
            result['ElementsDone'] += 1
        if status not in ['Done', 'Available']:
            agent = element['ChildQueueUrl']
            if agent not in result['ActiveAgents']:
                result['ActiveAgents'][agent] = 0
            result['ActiveAgents'][agent] += 1
    return result

def changePriority(requestName, priority, wmstatUrl = None):
    """
    Changes the priority that's stored in the workload.
    Takes the current priority stored in the workload and adds
    to it the input priority value. 
    
    """
    request = requestDetails(requestName)
    # change in Oracle
    newPrior = int(priority)
    ChangeState.changeRequestPriority(requestName, newPrior)
    # change in workload (spec)
    helper = loadWorkload(request)
    helper.data.request.priority = newPrior
    saveWorkload(helper, request['RequestWorkflow'], wmstatUrl)
    # change priority in CouchDB
    couchDb = Database(request["CouchWorkloadDBName"], request["CouchURL"])
    fields = {"RequestPriority": newPrior}
    couchDb.updateDocument(requestName, "ReqMgr", "updaterequest", fields=fields)
    # push the change to the WorkQueue
    response = ProdManagement.getProdMgr(requestName)
    if response == [] or response[0] is None or response[0] == "":
        # Request must not be assigned yet, we are safe here
        return
    workqueue = WorkQueue.WorkQueue(response[0])
    workqueue.updatePriority(requestName, priority)
    return

def main():
    """
    _main_
    """
    usage = "Usage: python %prog -w workflow"
    parser = OptionParser(usage = usage)
    parser.add_option('-w', '--workflow', help = 'Workflow name in ReqMgr', dest = 'wf')
    (options, args) = parser.parse_args()
    if not options.wf:
        parser.error('You must provide a workflow name')
        sys.exit(1)

    couchUrl = "https://cmsweb.cern.ch/couchdb"
    database = "acdcserver"
    failures = {}
    svc = Database(database, couchUrl)

    result = svc.loadView("ACDC", "byCollectionName", {'key' : options.wf, 'include_docs' : True, 'reduce' : False})
    print "Found %i failures/rows in total." % len(result["rows"])
    for entry in result["rows"]:
        if entry['doc']['fileset_name'] in failures:
            failures[entry['doc']['fileset_name']] += 1
        else:
            failures[entry['doc']['fileset_name']] = 1
    pprint(failures)
    print "\nDone!"

    def testSlashInDBName(self):
        """
        Slashes are a valid character in a database name, and are useful as it
        creates a directory strucutre for the couch data files.
        """
        db_name = 'wmcore/unittests'
        try:
            self.server.deleteDatabase(db_name)
        except:
            # Ignore this - the database shouldn't already exist
            pass

        db = self.server.createDatabase(db_name)
        info = db.info()
        assert info['db_name'] == db_name

        db_name = 'wmcore/unittests'
        db = self.server.connectDatabase(db_name)
        info = db.info()
        assert info['db_name'] == db_name

        db = Database(db_name, url = os.environ["COUCHURL"])
        info = db.info()
        assert info['db_name'] == db_name

        self.server.deleteDatabase(db_name)

def checkForMissingFiles(options):
    #Initialize stuff
    phedexAPI = PhEDEx({'cachepath' : options.cachepath})
    acdcCouch = Database('wmagent_acdc', options.acdcUrl)

    #Let's get the IDs of the ACDC documents for the task/request/group/user
    array = [options.group, options.user, options.request, options.task]
    result = acdcCouch.loadView('ACDC', 'owner_coll_fileset_docs', {'reduce' : False}, [array])

    documentsIDs = [x['id'] for x in result['rows']]
    
    badFiles = {}

    #Go through the documents
    for docID in documentsIDs:
        doc = acdcCouch.document(docID)

        #Are we going to change this doc? Better back it up
        if options.change:
            backupFile = os.open(os.path.join(options.backup, "%s.bkp" % doc["_id"]), 'w')
            json.dump(doc, backupFile)
            backupFile.close()

        #Go through the files
        files = doc["files"]
        for inputFile in files:

            #Use PhEDEx API to get site based on the SE
            se = files[inputFile]["locations"][0]
            siteLocation = phedexAPI.getBestNodeName(se)

            #Now get the PFN
            pfnDict = phedexAPI.getPFN(siteLocation, inputFile)
            inputPfn = pfnDict[(siteLocation, inputFile)]

            #Run lcg-ls commands and see what we get
            command = 'lcg-ls -b -D srmv2 --srm-timeout 60 %s' % inputPfn
            
            commandList = shlex.split(command)
            try:
                (stdout, stderr, exitCode) = runCommand(commandList, False, 70)
            except Exception, ex:
                exitCode = 99999
                stdout = ''
                stderr = str(ex)
            
            if exitCode:
                #Something went wrong with the command
                #Mark the file as bad
                if docID not in badFiles:
                    badFiles[docID] = []
                badFiles[docID].append(inputFile)
                print 'File %s is thought to be bad' % inputFile
                print 'Command was %s' % command
                print 'Return code was %i' % exitCode
                print 'Stdout was %s' % stdout
                print 'Stderr was %s' % stderr

def getSplitting(requestName):
	reqmgrUrl='https://cmsweb.cern.ch/reqmgr/reqMgr/'
	reqmgr = RequestManager(dict = {'endpoint' : reqmgrUrl})
	result = reqmgr.getRequest(requestName)
	workloadDB = Database(result['CouchWorkloadDBName'], result['CouchURL'])
	workloadPickle = workloadDB.getAttachment(requestName, 'spec')
	spec = pickle.loads(workloadPickle)
	workload = WMWorkloadHelper(spec)
	params = workload.getTopLevelTask()[0].jobSplittingParameters()
	algo = workload.getTopLevelTask()[0].jobSplittingAlgorithm()
	return params['events_per_job']

def main():
    demPolicy = Block()
    reqmgr = RequestManager(dict = {'endpoint' : 'https://cmsweb.cern.ch/reqmgr/reqMgr'})
    result = reqmgr.getRequest('pdmvserv_HIG-Summer12DR53X-01392_T1_ES_PIC_MSS_1_v0__130724_063344_7207')
    workloadDB = Database(result['CouchWorkloadDBName'], result['CouchURL'])
    workloadPickle = workloadDB.getAttachment('pdmvserv_HIG-Summer12DR53X-01392_T1_ES_PIC_MSS_1_v0__130724_063344_7207', 'spec')
    spec = pickle.loads(workloadPickle)
    workload = WMWorkloadHelper(spec)
    x,y = demPolicy(wmspec = workload, task = workload.getTopLevelTask()[0])
    print x
    print y

class RESTSink(object):
    """
    Alert sink for posting alerts to a REST server.
    The class acts as a REST client.

    """
    def __init__(self, config):
        # configuration values:
        #     'uri' attribute (URL of the REST server and resource name)
        #         in case of CouchDB, the resource name is the database name
        #         http://servername:port/databaseName
        self.config = config
        logging.info("Instantiating ...")

        # the class currently relies only on 1 REST server possibility - the
        # CouchDB server. as explained above, .database will be replaced by
        # .connection if both a generic REST server as well as CouchDB are to
        # be talked to
        split = self.config.uri.rfind('/')
        dbName = self.config.uri[split + 1:] # get last item of URI - database name
        url = self.config.uri[:split]
        # as opposed to CouchSink, here it's assumed the resource (the database name)
        # does exist, fail here otherwise
        # this check / rest of the constructed may be revised for
        #     general REST server
        server = CouchServer(url)
        databases = server.listDatabases()
        # there needs to be this database created upfront and also
        # couchapp associated with it installed, if it's there, fail
        if dbName not in databases:
            raise Exception("REST URI: %s (DB name: %s) does not exist." %
                            (self.config.uri, dbName))
        self._database = Database(dbName, url)
        logging.info("Initialized.")


    def send(self, alerts):
        """
        Send a list of alerts to a REST server.

        """
        for a in alerts:
            doc = Document(None, a)
            self._database.queue(doc)
        # two options here: either to call commit on the couch myself
        # or leave the alerts buffered in the Database queue which means
        # the .commit() would be called automatically if size is exceeded
        # 1st option:
        retVal = self._database.commit()
        logging.debug("Stored %s alerts to REST resource, retVals: %s" % (len(alerts), retVal))
        return retVal

def main():
    if "WMAGENT_CONFIG" not in os.environ:
        os.environ["WMAGENT_CONFIG"] = '/data/srv/wmagent/current/config/wmagent/config.py'
    myThread = threading.currentThread()
    connectToDB()
    formatter = DBFormatter(logging, myThread.dbi)
    limboFiles = formatter.formatDict(myThread.dbi.processData("""SELECT dbsbuffer_workflow.name, dbsbuffer_file.lfn
                                                                  FROM dbsbuffer_file
                                                                  INNER JOIN dbsbuffer_workflow ON
                                                                      dbsbuffer_file.workflow = dbsbuffer_workflow.id
                                                                  LEFT OUTER JOIN dbsbuffer_block ON
                                                                      dbsbuffer_file.block_id = dbsbuffer_block.id
                                                                  WHERE dbsbuffer_file.status = 'READY' AND
                                                                        dbsbuffer_block.id is NULL"""))
    if not limboFiles:
        print "There are no bad files to fix"
        return
    for entry in limboFiles:
        data = Database('wmagent_jobdump/fwjrs', 'http://%s:5984' % socket.gethostname())
        result = data.loadView('FWJRDump', 'jobsByOutputLFN', {'include_docs' : True},
                      [[entry['name'], entry['lfn']]])['rows']
        if result:
            result = result[0]
            fwjr = result['doc']['fwjr']
            for step in fwjr['steps']:
                if step == 'cmsRun1':
                    stepInfo = fwjr['steps'][step]
                    site = stepInfo['site']
                    break
        else:
            print "Could not find location for %s" % entry['lfn']
            continue
        se = myThread.dbi.processData("""SELECT wmbs_location_senames.se_name FROM
                                       wmbs_location_senames
                                       INNER JOIN wmbs_location ON
                                          wmbs_location.id = wmbs_location_senames.location
                                       WHERE wmbs_location.site_name = '%s'""" % site)
        se = formatter.formatDict(se)[0]
        insertQuery = """INSERT INTO dbsbuffer_location (se_name)
               SELECT '%s' AS se_name FROM DUAL WHERE NOT EXISTS
                (SELECT se_name FROM dbsbuffer_location WHERE se_name = '%s')""" % (se['se_name'], se['se_name'])
        myThread.dbi.processData(insertQuery)
        updateQuery = """INSERT INTO dbsbuffer_file_location (filename, location)
                           SELECT df.id, dl.id
                           FROM dbsbuffer_file df,  dbsbuffer_location dl
                           WHERE df.lfn = '%s'
                           AND dl.se_name = '%s'""" % (entry['lfn'], se['se_name'])
        myThread.dbi.processData(updateQuery)
        updateQuery = """UPDATE dbsbuffer_file SET status = 'NOTUPLOADED' WHERE lfn = '%s'""" % entry['lfn']
        myThread.dbi.processData(updateQuery)

def main():
    if len(sys.argv) != 2:
        print "Usage:"
        print "python CheckWorkQueueElements.py <workflowName>"
        sys.exit(0)

    workflow = sys.argv[1]
    x = Database('workqueue', 'https://cmsweb.cern.ch/couchdb')
    y = x.loadView('WorkQueue', 'elementsByParent', {'include_docs' : True}, [workflow])
    for entry in y['rows']:
        doc = entry['doc']
        element = doc['WMCore.WorkQueue.DataStructs.WorkQueueElement.WorkQueueElement']
        if element['Status'] != 'Done':
            print 'Element: %s is %s in %s' % (doc['_id'], element['Status'], element['ChildQueueUrl'])

def main():
    sum = 0
    x = Database('workqueue', 'http://vocms201.cern.ch:5984')
    y = x.loadView('WorkQueue', 'availableByPriority', {'include_docs' : True})
    loadDistribution = {}
    for entry in y['rows']:
        doc = entry['doc']
        element = doc['WMCore.WorkQueue.DataStructs.WorkQueueElement.WorkQueueElement']
        key = frozenset(element['SiteWhitelist'])
        if key not in loadDistribution:
            loadDistribution[key] = 0
        loadDistribution[key] += element['Jobs']
    for site, jobs in loadDistribution.items():
        print "Site list %s has %d jobs" % (str(site), jobs)

def main():
    print "Looking for problematic inbox elements..."
    problemRequests = getProblematicRequests()
    print "Found %d bad elements:" % len(problemRequests)
    if not problemRequests:
        print "Nothing to fix, contact a developer if the problem persists..."
        return 0
    for request in problemRequests:
        print request["RequestName"]
    var = raw_input("Can we close these for new data in inbox elements: Y/N\n")
    if var == "Y":
        print "Updating them in global inbox, you need a WMAgent proxy for this."
        inboxDB = Database('workqueue_inbox', 'https://cmsweb.cern.ch/couchdb')
        for request in problemRequests:
            inboxDB.document(request._id)
            inboxDB.updateDocument(request._id, 'WorkQueue', 'in-place', fields={'OpenForNewData': false})
        print "Done with the deletions, this should fix the problem."
        return 0
    else:
        var = raw_input("Then can we delete these inbox elements: Y/N\n")
        if var == "Y":
            print "Deleting them from the global inbox, you need a WMAgent proxy for this."
            inboxDB = Database('workqueue_inbox', 'https://cmsweb.cern.ch/couchdb')
            for request in problemRequests:
                inboxDB.delete_doc(request._id, request.rev)
            print "Done with the deletions, this should fix the problem."
            return 0
        else:
            print "Doing nothing as you commanded..."
        return 0