diff --git a/SAS/ResourceAssignment/ResourceAssignmentDatabase/radb.py b/SAS/ResourceAssignment/ResourceAssignmentDatabase/radb.py index f974a13f872d471726da1a16da22239204a4345e..a3833bc4b622442d632438eb57904f99b20cf603 100644 --- a/SAS/ResourceAssignment/ResourceAssignmentDatabase/radb.py +++ b/SAS/ResourceAssignment/ResourceAssignmentDatabase/radb.py @@ -155,7 +155,10 @@ class RADatabase: raise KeyError('No such status: %s. Valid values are: %s' % (status_name, ', '.join(self.getResourceClaimStatusNames()))) - def getTasks(self, lower_bound=None, upper_bound=None, task_ids=None, task_status=None, task_type=None): + def getTasks(self, lower_bound=None, upper_bound=None, task_ids=None, task_status=None, task_type=None, mom_ids=None, otdb_ids=None): + if len([x for x in [task_ids, mom_ids, otdb_ids] if x != None]) > 1: + raise KeyError("Provide either task_ids or mom_ids or otdb_ids, not multiple kinds.") + query = '''SELECT * from resource_allocation.task_view''' conditions = [] @@ -173,10 +176,26 @@ class RADatabase: if isinstance(task_ids, int): # just a single id conditions.append('id = %s') qargs.append(task_ids) - else: #assume a list/enumerable of id's + elif task_ids: #assume a list/enumerable of id's conditions.append('id in %s') qargs.append(tuple(task_ids)) + if mom_ids is not None: + if isinstance(mom_ids, int): # just a single id + conditions.append('mom_id = %s') + qargs.append(mom_ids) + elif mom_ids: #assume a list/enumerable of id's + conditions.append('mom_id in %s') + qargs.append(tuple(mom_ids)) + + if otdb_ids is not None: + if isinstance(otdb_ids, int): # just a single id + conditions.append('otdb_id = %s') + qargs.append(otdb_ids) + elif otdb_ids: #assume a list/enumerable of id's + conditions.append('otdb_id in %s') + qargs.append(tuple(otdb_ids)) + task_status, task_type = self._convertTaskTypeAndStatusToIds(task_status, task_type) if task_status is not None: diff --git a/SAS/ResourceAssignment/ResourceAssignmentService/rpc.py b/SAS/ResourceAssignment/ResourceAssignmentService/rpc.py index 7f8e3f22e7b6b98cf5cfd7a2203a9421fb62af28..44834eed42e913ff20f0679502f65d067ea87c1a 100644 --- a/SAS/ResourceAssignment/ResourceAssignmentService/rpc.py +++ b/SAS/ResourceAssignment/ResourceAssignmentService/rpc.py @@ -178,8 +178,8 @@ class RARPC(RPCWrapper): otdb_id=otdb_id, status=status) - def getTasks(self, lower_bound=None, upper_bound=None, task_ids=None, task_status=None, task_type=None): - tasks = self.rpc('GetTasks', lower_bound=lower_bound, upper_bound=upper_bound, task_ids=task_ids, task_status=task_status, task_type=task_type) + def getTasks(self, lower_bound=None, upper_bound=None, task_ids=None, task_status=None, task_type=None, mom_ids=None, otdb_ids=None): + tasks = self.rpc('GetTasks', lower_bound=lower_bound, upper_bound=upper_bound, task_ids=task_ids, task_status=task_status, task_type=task_type, mom_ids=mom_ids, otdb_ids=otdb_ids) for task in tasks: task['starttime'] = task['starttime'].datetime() task['endtime'] = task['endtime'].datetime() diff --git a/SAS/ResourceAssignment/ResourceAssignmentService/service.py b/SAS/ResourceAssignment/ResourceAssignmentService/service.py index 234019ae0a1b66e32b4292bff9e4e48da02f14de..7e01393bee29ff7c8f2d10f8ce52a50da8489286 100644 --- a/SAS/ResourceAssignment/ResourceAssignmentService/service.py +++ b/SAS/ResourceAssignment/ResourceAssignmentService/service.py @@ -215,7 +215,9 @@ class RADBHandler(MessageHandlerInterface): upper_bound=kwargs.get('upper_bound').datetime() if kwargs.get('upper_bound') else None, task_ids=kwargs.get('task_ids'), task_status=kwargs.get('task_status'), - task_type=kwargs.get('task_type')) + task_type=kwargs.get('task_type'), + mom_ids=kwargs.get('mom_ids'), + otdb_ids=kwargs.get('otdb_ids')) def _getTask(self, **kwargs): logger.info('GetTask: %s' % dict({k:v for k,v in kwargs.items() if v != None}))