class OperateCmd(QueryCmd): """An operation command will query the database and run an arbitrary operation on each object returned.""" # the default set of options for all operation commands options = [ CmdLineTool.StrListOption("-c", "--cols", help="list of columns that will be " "printed for each each row."), CmdLineTool.StrListOption("-s", "--sortby", help="sort the results in a specified " "order before performing any " "operations."), CmdLineTool.BooleanOption("--ns", "--nosort", dest="nosort", help="do not sort the results"), CmdLineTool.IntOption ("--limit", help="only operate on the first l rows"), ] + CmdLineTool.CmdLineTool.options def __init__(self, table, *args, **kwargs): # call the super super(OperateCmd, self).__init__(table, *args, **kwargs) # if a list of objects is provided at the command line, then # subclasses can fill in this value. Setting this will prevent # the database from being contacted self.objects = [] def parseArgs(self, *args, **kwargs): result = super(OperateCmd, self).parseArgs(*args, **kwargs) # make sure a where exists if not self.where: raise OptionParser.HelpRequest(self.getHelpStr()) # make sure a value is set for options we are not supporting, but # that the super expects to exist. for name,default in (("distinct", None), ("delimiter", ' '), ("raw", False), ("timefmt", None), ("noformat", False), ("noheader", False)): self.opts.__dict__.setdefault(name, default) return result def opendb(self): """We don't need to open the database if the objects member is set.""" if not self.objects: super(OperateCmd, self).opendb() def getNewObjs(self, num, curr, objs, distinct): """Return at most 'num' new objects for the operationLoop. The objects returned are tq Operation Job or Task objects. By default the mrd cursor is checked for waiting job objects, but a list can be provided which will be modified.""" # a list of all the operation objects we are going to return newobjs = [] # loop until we have enough objects while len(newobjs) < num and curr < len(objs): obj = objs[curr] curr += 1 # check if this object is distinct if not self.isDistinct(obj, distinct): if self.parent.opts.debug: print 'already worked on' continue # skip over this object if we don't have permission if not (self.parent.force or obj.hasPermission()): msg = "permission denied: not job owner" if self.parent.color: msg = terminal.TerminalColor('red').colorStr(msg) print self.formatter.format(obj), msg continue # we passed all the tests, so add it to the list newobjs.append(obj) return newobjs,curr def processResult(self, qresult): """Process and print all the objects returned from the query.""" # if we had no results, then do nothing if not qresult or \ (self.opts.limit is not None and self.opts.limit < 0): return # make sure the query result doesn't cache objects after they are # created, as we will only need them once. if isinstance(qresult, Database.QueryResult): qresult.cacheObjects = False file = sys.stdout # set the widths of our formatter so everything is neatly spaced if not self.opts.noformat: self.formatter.setWidths(qresult) # print a header if not self.opts.noheader: print >>file, self.formatter.header() # keep track of which objects are distinct distinct = {} # our current list of objects that have had the operation started, # but the operation has not finished. Typically because the dispatcher # is busy reading from it and would block the process. queue = [] # keep track of how many objects we've worked on workedon = 0 # time we last created an object (only useful when pausing between # operations). lastcreate = 0 # make sure threads is set to at least one if self.parent.threads < 1: self.parent.threads = 1 # setup a blocking flag if not self.parent.yes or self.parent.threads <= 1: self.blocking = True else: self.blocking = False # keep track of our current index into the query result. curr = 0 try: # start operating on the objects while True: # if we are running in parallel mode (i.e. non-blocking) we # want to have as many concurrent connections as possible. # Thus make sure the queue size is always the size of # self.threads, or at least as big as it can be. newobjs = [] # check if we need to pause inbetween operations if self.parent.pause > 0: # how much time has passed since our last object create elapsed = self.parent.pause - (time.time() - lastcreate) # if we don't have any objects on the queue, and 'pause' # seconds hasn't elapsed if not queue and elapsed > 0: time.sleep(elapsed) elapsed = 0 # add another item to the queue if elapsed <= 0 and len(queue) < self.parent.threads: newobjs,curr = self.getNewObjs(1, curr, qresult, distinct) lastcreate = time.time() else: # get however many more we can take on now newobjs,curr = self.getNewObjs( self.parent.threads - len(queue), curr, qresult, distinct) # if we have nothing left, then we're all done! if not (queue or newobjs) or \ (not queue and \ self.opts.limit not in [None, 0] and \ self.opts.limit <= workedon): break # start the new operations now. If we are in blocking # mode, (self.yesToAll == False or self.threads <= 1), # then the entire operation will be done here. If we are # in parallel mode (i.e. non-blocking), then the connect # will initialize here, raise a WouldBlock exception, and # be put onto the queue. for obj in newobjs: try: # work on this object self.processObject(obj) except Sockets.WouldBlock: # if this would block, then add it to the queue queue.append(obj) except Sockets.SocketError, err: # catch any other socket related error obj.post(self, str(err), error=True) obj.close() except EngineClient.EngineClientError, err: obj.post(self, str(err), error=True) obj.close() except:
class ChangeCrewsCmd(JobOperationCmd): usage = "chcrews" description = """ Change the crews of one or more jobs. The crews of a job determine which blades the job's tasks will potentially run on. Crews can be added or removed from a job's existing list, or replaced entirely with a new list of crews. By default the user is prompted before each job has its crews changed, unless the --yes flag is set. Also, you can only change the crews of your own job, unless the --force flag is set. """ examples = """ Syntax to change the crews of jobs: > tq chcrews SEARCH_CLAUSE [-k|--crews|-a|-r] crew1[,crew2,..] > tq chcrews jid [jid2 jid3 ...] [-k|--crews|-a|-r] crew1[,crew2,..] Examples: change the crews of all of joni's jobs to lighting > tq chcrews user=joni -k lighting add animation to all of joni's jobs > tq chcrews user=joni -a animation remove the animation crew from a specific job > tq chcrews 10795593 -r animation """ # add option to get new crews options = [ CmdLineTool.StrListOption("-a", "--add", help="a comma delimited list of " "crews that will be added to " "each job's current list."), CmdLineTool.StrListOption("-k", "--crews", help="a comma delimited list of " "crews to set as the job's " "crews."), CmdLineTool.StrListOption("-r", "--remove", help="a comma delimited list of " "crews that will be removed from " "each job's current list."), ] + JobOperationCmd.options def parseArgs(self, *args, **kwargs): """Make sure we grab the current crews if we need to add or remove.""" result = super(ChangeCrewsCmd, self).parseArgs(*args, **kwargs) # make sure crews were provided if not (self.opts.add or self.opts.remove or self.opts.crews): raise tq.TqError, "no crews provided" # if the add or remove options are set, then force a database call if self.objects and (self.opts.add or self.opts.remove): phrases = ['jobid=%d' % job.jobid for job in self.objects] self.where = " or ".join(phrases) # reset the objects list so we do a database call self.objects = [] return result def pre_execute(self): super(ChangeCrewsCmd, self).pre_execute() # add the crews if we are modifying if self.opts.add or self.opts.remove: self.members.append("Job.crews") def processObject(self, obj): """Operate on the provided job object.""" # ask the user if we should continue if not obj.pre(self, "Change the crews of the job?"): return if self.opts.add or self.opts.remove: # make sure everything is unique origcrews = listutil.getUnion(obj.crews) # make a copy so we can add/remove crews newcrews = list(origcrews) for crew in (self.opts.add or self.opts.remove): if self.opts.add: newcrews.append(crew) else: if crew in origcrews: newcrews.remove(crew) else: obj.post( self, "crew %s is not in crews %s" % (crew, origcrews)) # make sure everything is unique again newcrews = listutil.getUnion(newcrews) if newcrews == origcrews: obj.post(self, "doesn't need to change") return crews = newcrews else: crews = listutil.getUnion(self.opts.crews) # try to run the operation query.chcrews(obj, crews=crews) obj.post(self, "crews changed")