addExperimentCli.py

#!/usr/bin/env python

from dm.aps_bss.api.apsBssApi import ApsBssApi
from dm.ds_web_service.api.experimentDsApi import ExperimentDsApi
from dm.ds_web_service.api.userDsApi import UserDsApi

from dm.common.exceptions.invalidRequest import InvalidRequest
from dm.common.utility.configurationManager import ConfigurationManager
from dsWebServiceSessionCli import DsWebServiceSessionCli

class AddExperimentCli(DsWebServiceSessionCli):
    def __init__(self):
        DsWebServiceSessionCli.__init__(self)
        configManager = ConfigurationManager.getInstance()
        self.allowedExperimentTypes = configManager.getAllowedExperimentTypes()
        allowedExperimentTypesHelp = ''
        if self.allowedExperimentTypes:
            allowedExperimentTypesHelp = ' Allowed types: %s' % self.allowedExperimentTypes
            

        self.addOption('', '--experiment', dest='experimentName', help='Experiment name.')
        self.addOption('', '--station', dest='stationName', help='Experiment station name, can also be set via DM_STATION_NAME environment variable.')
        self.addOption('', '--type', dest='typeName', help='Experiment type name.%s' % allowedExperimentTypesHelp)
        self.addOption('', '--type-id', dest='typeId', help='Experiment type id (may be given instead of type name; type id is ignored if both type name and id are provided).')
        self.addOption('', '--description', dest='description', help='Experiment description.')
        self.addOption('', '--start-date', dest='startDate', help='Experiment start date in format DD-MMM-YY.')
        self.addOption('', '--end-date', dest='endDate', help='Experiment end date in format DD-MMM-YY.')
        self.addOption('', '--users', dest='users', help='Comma-separated list of DM usernames to be added to the new experiment as users.')
        self.addOption('', '--proposal-id', dest='proposalId', help='Beamline proposal id. If specified, all users listed on the proposal will be added to the new experiment.')
        self.addOption('', '--run', dest='runName', help='Run name. If not specified, current run name is assumed for beamline proposal.')
        self.addOption('', '--bss-login-file', dest='bssLoginFile', help='BSS login file. Login file may also be specified via environment variable DM_BSS_LOGIN_FILE.')

    def checkArgs(self):
        if self.options.experimentName is None:
            raise InvalidRequest('Experiment name must be provided.')
        if self.getTypeName() is None:
            raise InvalidRequest('Experiment type name must be provided.')
        if self.getStationName() is None:
            raise InvalidRequest('Experiment station name must be provided.')
        # If allowed experiment types is not set, there are no restrictions
        if self.allowedExperimentTypes:
            if self.getTypeName() not in self.allowedExperimentTypes.split(','):
                raise InvalidRequest('Experiment type %s is not allowed on this station. Allowed types are: %s.' % (self.getTypeName(), self.allowedExperimentTypes))

    def getExperimentName(self):
        return self.options.experimentName

    def getStationName(self):
        stationName = self.options.stationName
        if not stationName:
            configManager = ConfigurationManager.getInstance()
            stationName = configManager.getStationName()
        return stationName

    def getTypeName(self):
        typeName = self.options.typeName
        if not typeName:
            if self.options.typeId:
                api = ExperimentDsApi(self.getLoginUsername(), self.getLoginPassword(), self.getServiceHost(), self.getServicePort(), self.getServiceProtocol())
                experimentType = api.getExperimentTypeById(typeId)
                typeName = experimentType.get('name')
                self.options.typeName = typeName
        return typeName

    def getDescription(self):
        return self.options.description

    def getStartDate(self):
        return self.options.startDate

    def getEndDate(self):
        return self.options.endDate

    def getProposalId(self):
        proposalId = self.options.proposalId
        if proposalId:
            proposalId = int(proposalId)
        return proposalId

    def getUsers(self):
        # Return list of users and beamline managers that can access data
        users = self.options.users
        if users:
            users = users.split(',')
        else:
            users = []
        beamlineManagers = ConfigurationManager.getInstance().getBeamlineManagers()
        if beamlineManagers:
            beamlineManagers = beamlineManagers.split(',')
        else:
            beamlineManagers = []
        # Remove duplicates by converting into set
        return list(set(users+beamlineManagers))

    def runCommand(self):
        self.parseArgs(usage="""
    dm-add-experiment 
        --experiment=EXPERIMENTNAME --station=STATIONNAME --type=TYPENAME|--type-id=TYPEID
        [--description=DESCRIPTION] 
        [--start-date=STARTDATE] 
        [--end-date=ENDDATE]
        [--users=USERS]
        [--proposal-id=PROPOSALID]
        [--run=RUNNAME]
        [--bss-login-file=BSSLOGINFILE]

Description:
    Add new experiment to the DM database. If list of users or proposal id is specified, this command will also add roles for all users listed on the proposal.
        """)
        self.checkArgs()
        dsExperimentApi = ExperimentDsApi(self.getLoginUsername(), self.getLoginPassword(), self.getServiceHost(), self.getServicePort(), self.getServiceProtocol())
        dsUserApi = UserDsApi(self.getLoginUsername(), self.getLoginPassword(), self.getServiceHost(), self.getServicePort(), self.getServiceProtocol())

        description = self.getDescription()
        proposalId = self.getProposalId()
        experimenters = []
        if proposalId:
            bssApi = ApsBssApi(loginFile=self.options.bssLoginFile)
            proposal = bssApi.getBeamlineProposal(proposalId=proposalId, runName=self.options.runName)
            experimenters = proposal.get('experimenters', [])
            if not description:
                description = '%s (Proposal id: %s)' % (proposal['title'], proposalId)


        users = self.getUsers()
        pis = []
        for experimenter in experimenters:
            badge = int(experimenter['badge'])
            if not badge:
                #print 'Skipping user %s due to invalid badge.' % lastName 
                continue
            username = 'd%s' % badge

            # Clasify user
            if experimenter.get('piFlag') == 'Y':
                if not pis.count(username):
                    pis.append(username)
                if users.count(username):
                    users.remove(username)
            else:
                if not users.count(username):
                    users.append(username)

        for username in users+pis:
            # Check that user exists
            dsUserApi.getUserByUsername(username)
           

        # Everything looks good, add experiment and users
        experiment = dsExperimentApi.addExperiment(self.getExperimentName(), self.getStationName(), self.getTypeName(), description, self.getStartDate(), self.getEndDate())

        # Add pis.
        experimentName = experiment['name']
        roleName = 'PI'
        for username in pis:
            dsUserApi.addUserExperimentRole(username, roleName, experimentName)
        roleName = 'User'
        for username in users:
            dsUserApi.addUserExperimentRole(username, roleName, experimentName)

        if len(users+pis):
            experiment = dsExperimentApi.getExperimentByName(experimentName)
        print experiment.getDisplayString(self.getDisplayKeys(), self.getDisplayFormat())

#######################################################################
# Run command.
if __name__ == '__main__':
    cli = AddExperimentCli()
    cli.run()