class Scheduler(object):
def __init__(self):
self.scheduler = GeventScheduler()
self.scheduler.add_listener(self.__scheduler_listener(),
EVENT_SCHEDULER_START | EVENT_SCHEDULER_SHUTDOWN
| EVENT_SCHEDULER_PAUSED | EVENT_SCHEDULER_RESUMED
| EVENT_JOB_ADDED | EVENT_JOB_REMOVED
| EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
self.id = 'controller'
self.app = None
def schedule_workflows(self, task_id, executable, workflow_ids, trigger):
"""
Schedules a workflow for execution
Args:
task_id (int): Id of the scheduled task
executable (func): A callable to execute must take in one argument -- a workflow id
workflow_ids (iterable(str)): An iterable of workflow ids
trigger (Trigger): The trigger to use for this scheduled task
"""
def execute(id_):
with self.app.app_context():
executable(id_)
for workflow_id in workflow_ids:
self.scheduler.add_job(execute, args=(workflow_id,),
id=construct_task_id(task_id, workflow_id),
trigger=trigger, replace_existing=True)
def get_all_scheduled_workflows(self):
"""
Gets all the scheduled workflows
Returns:
(dict{str: list[str]}) A dict of task_id to workflow execution ids
"""
tasks = {}
for job in self.scheduler.get_jobs():
task, workflow_execution_id = split_task_id(job.id)
if task not in tasks:
tasks[task] = [workflow_execution_id]
else:
tasks[task].append(workflow_execution_id)
return tasks
def get_scheduled_workflows(self, task_id):
"""
Gets all the scheduled worfklows for a given task id
Args:
task_id (str): The task id
Returns:
(list[str]) A list fo workflow execution id associated with this task id
"""
tasks = []
for job in self.scheduler.get_jobs():
task, workflow_execution_id = split_task_id(job.id)
if task == task_id:
tasks.append(workflow_execution_id)
return tasks
def update_workflows(self, task_id, trigger):
"""
Updates the workflows for a given task id to use a different trigger
Args:
task_id (str|int): The task id to update
trigger (Trigger): The new trigger to use
"""
existing_tasks = {construct_task_id(task_id, workflow_execution_id) for workflow_execution_id in
self.get_scheduled_workflows(task_id)}
for job_id in existing_tasks:
self.scheduler.reschedule_job(job_id=job_id, trigger=trigger)
def unschedule_workflows(self, task_id, workflow_execution_ids):
"""
Unschedules a workflow
Args:
task_id (str|int): The task ID to unschedule
workflow_execution_ids (list[str]): The list of workflow execution IDs to update
"""
for workflow_execution_id in workflow_execution_ids:
try:
self.scheduler.remove_job(construct_task_id(task_id, workflow_execution_id))
except JobLookupError:
logger.warning('Cannot delete task {}. '
'No task found in scheduler'.format(construct_task_id(task_id, workflow_execution_id)))
def start(self):
"""Starts the scheduler for active execution. This function must be called before any workflows are executed.
Returns:
The state of the scheduler if successful, error message if scheduler is in "stopped" state.
"""
if self.scheduler.state == STATE_STOPPED:
logger.info('Starting scheduler')
self.scheduler.start()
else:
logger.warning('Cannot start scheduler. Scheduler is already running or is paused')
return "Scheduler already running."
return self.scheduler.state
def stop(self, wait=True):
"""Stops active execution.
Args:
wait (bool, optional): Boolean to synchronously or asynchronously wait for the scheduler to shutdown.
Default is True.
Returns:
The state of the scheduler if successful, error message if scheduler is already in "stopped" state.
"""
if self.scheduler.state != STATE_STOPPED:
logger.info('Stopping scheduler')
self.scheduler.shutdown(wait=wait)
else:
logger.warning('Cannot stop scheduler. Scheduler is already stopped')
return "Scheduler already stopped."
return self.scheduler.state
def pause(self):
"""Pauses active execution.
Returns:
The state of the scheduler if successful, error message if scheduler is not in the "running" state.
"""
if self.scheduler.state == STATE_RUNNING:
logger.info('Pausing scheduler')
self.scheduler.pause()
elif self.scheduler.state == STATE_PAUSED:
logger.warning('Cannot pause scheduler. Scheduler is already paused')
return "Scheduler already paused."
elif self.scheduler.state == STATE_STOPPED:
logger.warning('Cannot pause scheduler. Scheduler is stopped')
return "Scheduler is in STOPPED state and cannot be paused."
return self.scheduler.state
def resume(self):
"""Resumes active execution.
Returns:
The state of the scheduler if successful, error message if scheduler is not in the "paused" state.
"""
if self.scheduler.state == STATE_PAUSED:
logger.info('Resuming scheduler')
self.scheduler.resume()
else:
logger.warning("Scheduler is not in PAUSED state and cannot be resumed.")
return "Scheduler is not in PAUSED state and cannot be resumed."
return self.scheduler.state
def pause_workflows(self, task_id, workflow_execution_ids):
"""
Pauses some workflows associated with a task
Args:
task_id (int|str): The id of the task to pause
workflow_execution_ids (list[str]): The list of workflow execution IDs to pause
"""
for workflow_execution_id in workflow_execution_ids:
job_id = construct_task_id(task_id, workflow_execution_id)
try:
self.scheduler.pause_job(job_id=job_id)
logger.info('Paused job {0}'.format(job_id))
except JobLookupError:
logger.warning('Cannot pause scheduled workflow {}. Workflow ID not found'.format(job_id))
def resume_workflows(self, task_id, workflow_execution_ids):
"""
Resumes some workflows associated with a task
Args:
task_id (int|str): The id of the task to pause
workflow_execution_ids (list[str]): The list of workflow execution IDs to resume
"""
for workflow_execution_id in workflow_execution_ids:
job_id = construct_task_id(task_id, workflow_execution_id)
try:
self.scheduler.resume_job(job_id=job_id)
logger.info('Resumed job {0}'.format(job_id))
except JobLookupError:
logger.warning('Cannot resume scheduled workflow {}. Workflow ID not found'.format(job_id))
def __scheduler_listener(self):
event_selector_map = {EVENT_SCHEDULER_START: WalkoffEvent.SchedulerStart,
EVENT_SCHEDULER_SHUTDOWN: WalkoffEvent.SchedulerShutdown,
EVENT_SCHEDULER_PAUSED: WalkoffEvent.SchedulerPaused,
EVENT_SCHEDULER_RESUMED: WalkoffEvent.SchedulerResumed,
EVENT_JOB_ADDED: WalkoffEvent.SchedulerJobAdded,
EVENT_JOB_REMOVED: WalkoffEvent.SchedulerJobRemoved,
EVENT_JOB_EXECUTED: WalkoffEvent.SchedulerJobExecuted,
EVENT_JOB_ERROR: WalkoffEvent.SchedulerJobError}
def event_selector(event):
try:
event = event_selector_map[event.code]
event.send(self)
except KeyError: # pragma: no cover
logger.error('Unknown event sent triggered in scheduler {}'.format(event))
return event_selector
class Controller(object):
def __init__(self, name="defaultController"):
self.name = name
self.workflows = {}
self.load_all_workflows_from_directory()
self.instances = {}
self.tree = None
self.eventlog = []
self.schedulerStatusListener = SchedulerStatusListener(self.eventlog)
self.jobStatusListener = JobStatusListener(self.eventlog)
self.jobExecutionListener = JobExecutionListener(self.eventlog)
self.scheduler = GeventScheduler()
self.scheduler.add_listener(
self.schedulerStatusListener.callback(self),
EVENT_SCHEDULER_START | EVENT_SCHEDULER_SHUTDOWN
| EVENT_SCHEDULER_PAUSED | EVENT_SCHEDULER_RESUMED)
self.scheduler.add_listener(self.jobStatusListener.callback(self),
EVENT_JOB_ADDED | EVENT_JOB_REMOVED)
self.scheduler.add_listener(self.jobExecutionListener.callback(self),
EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
self.ancestry = [self.name]
# MULTIPROCESSING
# self.pool = multiprocessing.Pool(processes=5)
# self.manager = multiprocessing.Manager()
# self.queue = self.manager.SimpleQueue()
def load_workflow_from_file(self,
path,
workflow_name,
name_override=None,
playbook_override=None):
self.tree = et.ElementTree(file=path)
playbook_name = playbook_override if playbook_override else os.path.splitext(
os.path.basename(path))[0]
for workflow in self.tree.iter(tag="workflow"):
current_workflow_name = workflow.get('name')
if current_workflow_name == workflow_name:
if name_override:
workflow_name = name_override
name = construct_workflow_name_key(playbook_name,
workflow_name)
key = _WorkflowKey(playbook_name, workflow_name)
self.workflows[key] = wf.Workflow(name=name,
workflowConfig=workflow,
parent_name=self.name,
filename=playbook_name)
break
else:
return False
self.addChildWorkflows()
self.addWorkflowScheduledJobs()
return True
def loadWorkflowsFromFile(self,
path,
name_override=None,
playbook_override=None):
self.tree = et.ElementTree(file=path)
playbook_name = playbook_override if playbook_override else os.path.splitext(
os.path.basename(path))[0]
for workflow in self.tree.iter(tag='workflow'):
workflow_name = name_override if name_override else workflow.get(
'name')
name = construct_workflow_name_key(playbook_name, workflow_name)
key = _WorkflowKey(playbook_name, workflow_name)
self.workflows[key] = wf.Workflow(name=name,
workflowConfig=workflow,
parent_name=self.name,
filename=playbook_name)
self.addChildWorkflows()
self.addWorkflowScheduledJobs()
def load_all_workflows_from_directory(self, path=config.workflowsPath):
for workflow in locate_workflows_in_directory(path):
self.loadWorkflowsFromFile(
os.path.join(config.workflowsPath, workflow))
def addChildWorkflows(self):
for workflow in self.workflows:
playbook_name = workflow.playbook
children = self.workflows[workflow].options.children
for child in children:
workflow_key = _WorkflowKey(
playbook_name,
extract_workflow_name(child, playbook_name=playbook_name))
if workflow_key in self.workflows:
children[child] = self.workflows[workflow_key]
def addWorkflowScheduledJobs(self):
for workflow in self.workflows:
if (self.workflows[workflow].options.enabled
and self.workflows[workflow].options.scheduler["autorun"]
== "true"):
schedule_type = self.workflows[workflow].options.scheduler[
"type"]
schedule = self.workflows[workflow].options.scheduler["args"]
self.scheduler.add_job(self.workflows[workflow].execute,
trigger=schedule_type,
replace_existing=True,
**schedule)
def create_workflow_from_template(self,
playbook_name,
workflow_name,
template_playbook='emptyWorkflow',
template_name='emptyWorkflow'):
path = '{0}{1}{2}.workflow'.format(config.templatesPath, sep,
template_playbook)
return self.load_workflow_from_file(path=path,
workflow_name=template_name,
name_override=workflow_name,
playbook_override=playbook_name)
def create_playbook_from_template(self,
playbook_name,
template_playbook='emptyWorkflow'):
#TODO: Need a handler for returning workflow key and status
path = '{0}{1}{2}.workflow'.format(config.templatesPath, sep,
template_playbook)
self.loadWorkflowsFromFile(path=path, playbook_override=playbook_name)
def removeWorkflow(self, playbook_name, workflow_name):
name = _WorkflowKey(playbook_name, workflow_name)
if name in self.workflows:
del self.workflows[name]
return True
return False
def remove_playbook(self, playbook_name):
for name in [
workflow for workflow in self.workflows
if workflow.playbook == playbook_name
]:
del self.workflows[name]
return True
return False
def get_all_workflows(self):
result = {}
for key in self.workflows.keys():
if key.playbook not in result:
result[key.playbook] = []
result[key.playbook].append(key.workflow)
return result
def is_workflow_registered(self, playbook_name, workflow_name):
return _WorkflowKey(playbook_name, workflow_name) in self.workflows
def is_playbook_registerd(self, playbook_name):
return any(workflow_key.playbook == playbook_name
for workflow_key in self.workflows)
def update_workflow_name(self, old_playbook, old_workflow, new_playbook,
new_workflow):
old_key = _WorkflowKey(old_playbook, old_workflow)
new_key = _WorkflowKey(new_playbook, new_workflow)
self.workflows[new_key] = self.workflows.pop(old_key)
self.workflows[new_key].name = construct_workflow_name_key(
new_playbook, new_workflow)
def update_playbook_name(self, old_playbook, new_playbook):
for key in [
name for name in self.workflows.keys()
if name.playbook == old_playbook
]:
self.update_workflow_name(old_playbook, key.workflow, new_playbook,
key.workflow)
# def executeWorkflowWorker(self):
#
# print("Thread " + str(os.getpid()) + " starting up...")
#
# while (True):
# while (self.queue.empty()):
# continue
# name,start,data = self.queue.get()
# print("Thread " + str(os.getpid()) + " received and executing workflow "+name)
# steps, instances = self.workflows[name].execute(start=start, data=data)
def executeWorkflow(self,
playbook_name,
workflow_name,
start="start",
data=None):
self.workflows[_WorkflowKey(playbook_name,
workflow_name)].execute(start=start,
data=data)
# print("Boss thread putting "+name+" workflow on queue...:")
# self.queue.put((name, start, data))
self.jobExecutionListener.execute_event_code(self, 'JobExecuted')
def get_workflow(self, playbook_name, workflow_name):
key = _WorkflowKey(playbook_name, workflow_name)
if key in self.workflows:
return self.workflows[key]
return None
def playbook_to_xml(self, playbook_name):
workflows = [
workflow for key, workflow in self.workflows.items()
if key.playbook == playbook_name
]
if workflows:
xml = et.Element("workflows")
for workflow in workflows:
xml.append(workflow.to_xml())
return xml
else:
return None
# Starts active execution
def start(self):
self.scheduler.start()
# Stops active execution
def stop(self, wait=True):
self.scheduler.shutdown(wait=wait)
# Pauses active execution
def pause(self):
self.scheduler.pause()
# Resumes active execution
def resume(self):
self.scheduler.resume()
# Pauses active execution of specific job
def pauseJob(self, job_id):
self.scheduler.pause_job(job_id=job_id)
# Resumes active execution of specific job
def resumeJob(self, job_id):
self.scheduler.resume_job(job_id=job_id)
# Returns jobs scheduled for active execution
def getScheduledJobs(self):
self.scheduler.get_jobs()
class W5Timer(object):
_instance_lock = threading.Lock()
def __init__(self):
self.scheduler = None
def __new__(cls, *args, **kwargs):
if not hasattr(W5Timer, "_instance"):
with W5Timer._instance_lock:
if not hasattr(W5Timer, "_instance"):
W5Timer._instance = object.__new__(cls)
return W5Timer._instance
def create_scheduler(self):
self.scheduler = GeventScheduler()
def start(self):
self.scheduler.start()
def shutdown(self):
self.scheduler.shutdown()
def pause(self, uuid):
self.scheduler.pause_job(uuid)
def pause_all(self):
self.scheduler.pause()
def resume(self, uuid):
self.scheduler.resume_job(uuid)
def resume_all(self):
self.scheduler.resume()
def remove_job(self, uuid):
self.scheduler.remove_job(uuid)
def get_jobs(self):
return self.scheduler.get_jobs()
def add_date(self, run_date=None, uuid=None, timer_uuid=None):
self.scheduler.add_job(
auto_execute,
'date',
run_date=run_date,
id=timer_uuid,
args=(uuid,)
)
def update_date(self, uuid, run_date=None):
self.scheduler.reschedule_job(
uuid,
trigger='date',
run_date=run_date
)
def add_interval(self, t, interval, uuid=None, timer_uuid=None, start_date=None, end_date=None, jitter=0):
if t == "seconds":
self.scheduler.add_job(
auto_execute,
'interval',
seconds=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter,
id=timer_uuid,
args=(uuid,)
)
elif t == "minutes":
self.scheduler.add_job(
auto_execute,
'interval',
minutes=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter,
id=timer_uuid,
args=(uuid,)
)
elif t == "hours":
self.scheduler.add_job(
auto_execute,
'interval',
hours=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter,
id=timer_uuid,
args=(uuid,)
)
elif t == "days":
self.scheduler.add_job(
auto_execute,
'interval',
days=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter,
id=timer_uuid,
args=(uuid,)
)
elif t == "weeks":
self.scheduler.add_job(
auto_execute,
'interval',
weeks=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter,
id=timer_uuid,
args=(uuid,)
)
def update_interval(self, uuid, t, interval, start_date=None, end_date=None, jitter=0):
if t == "seconds":
self.scheduler.reschedule_job(
uuid,
trigger="interval",
seconds=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter
)
elif t == "minutes":
self.scheduler.reschedule_job(
uuid,
trigger="interval",
minutes=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter
)
elif t == "hours":
self.scheduler.reschedule_job(
uuid,
trigger="interval",
hours=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter
)
elif t == "days":
self.scheduler.reschedule_job(
uuid,
trigger="interval",
days=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter
)
elif t == "weeks":
self.scheduler.reschedule_job(
uuid,
trigger="interval",
weeks=interval,
start_date=start_date,
end_date=end_date,
jitter=jitter
)
def add_cron(self, cron, uuid=None, timer_uuid=None, start_date=None, end_date=None, jitter=0):
self.scheduler.add_job(
auto_execute,
CronTrigger.from_crontab(cron),
start_date=start_date,
end_date=end_date,
jitter=jitter,
id=timer_uuid,
args=(uuid,)
)
def update_cron(self, uuid, cron, start_date=None, end_date=None, jitter=0):
values = cron.split()
if len(values) != 5:
raise ValueError('Wrong number of fields; got {}, expected 5'.format(len(values)))
self.scheduler.reschedule_job(
uuid,
None,
"cron",
minute=values[0],
hour=values[1],
day=values[2],
month=values[3],
day_of_week=values[4],
start_date=start_date,
end_date=end_date,
jitter=jitter
)
class Controller(object):
def __init__(self,
name='defaultController',
workflows_path=core.config.paths.workflows_path):
"""Initializes a Controller object.
Args:
name (str, optional): Name for the controller.
workflows_path (str, optional): Path to the workflows.
"""
self.name = name
self.workflows = {}
self.load_all_workflows_from_directory(path=workflows_path)
self.instances = {}
self.tree = None
self.scheduler = GeventScheduler()
self.scheduler.add_listener(
self.__scheduler_listener(),
EVENT_SCHEDULER_START | EVENT_SCHEDULER_SHUTDOWN
| EVENT_SCHEDULER_PAUSED | EVENT_SCHEDULER_RESUMED
| EVENT_JOB_ADDED | EVENT_JOB_REMOVED
| EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
self.ancestry = [self.name]
self.paused_workflows = {}
def reconstruct_ancestry(self):
"""Reconstructs the ancestry list field of a workflow in case it changes.
"""
for key in self.workflows:
self.workflows[key].reconstruct_ancestry(self.ancestry)
def load_workflow_from_file(self,
path,
workflow_name,
name_override=None,
playbook_override=None):
"""Loads a workflow from a file.
Args:
path (str): Path to the workflow.
workflow_name (str): Name of the workflow to load.
name_override (str, optional): Name that the workflow should be changed to.
playbook_override (str, optional): Name that the playbook should be changed to.
Returns:
True on success, False otherwise.
"""
self.tree = cElementTree.ElementTree(file=path)
playbook_name = playbook_override if playbook_override else os.path.splitext(
os.path.basename(path))[0]
for workflow in self.tree.iter(tag='workflow'):
current_workflow_name = workflow.get('name')
if current_workflow_name == workflow_name:
if name_override:
workflow_name = name_override
name = construct_workflow_name_key(playbook_name,
workflow_name)
key = _WorkflowKey(playbook_name, workflow_name)
self.workflows[key] = wf.Workflow(name=name,
xml=workflow,
parent_name=self.name,
playbook_name=playbook_name)
logger.info('Adding workflow {0} to controller'.format(name))
break
else:
logger.warning(
'Workflow {0} not found in playbook {0}. Cannot load.'.format(
workflow_name, playbook_name))
return False
self.add_child_workflows()
self.add_workflow_scheduled_jobs()
return True
def load_workflows_from_file(self,
path,
name_override=None,
playbook_override=None):
"""Loads multiple workloads from a file.
Args:
path (str): Path to the workflow.
name_override (str, optional): Name that the workflow should be changed to.
playbook_override (str, optional): Name that the playbook should be changed to.
"""
self.tree = cElementTree.ElementTree(file=path)
playbook_name = playbook_override if playbook_override else os.path.splitext(
os.path.basename(path))[0]
for workflow in self.tree.iter(tag='workflow'):
workflow_name = name_override if name_override else workflow.get(
'name')
name = construct_workflow_name_key(playbook_name, workflow_name)
key = _WorkflowKey(playbook_name, workflow_name)
self.workflows[key] = wf.Workflow(name=name,
xml=workflow,
parent_name=self.name,
playbook_name=playbook_name)
logger.info('Adding workflow {0} to controller'.format(name))
self.add_child_workflows()
self.add_workflow_scheduled_jobs()
def load_all_workflows_from_directory(self, path=None):
"""Loads all workflows from a directory.
Args:
path (str, optional): Path to the directory to load from. Defaults to the configuration workflows_path.
"""
if path is None:
path = core.config.paths.workflows_path
for workflow in locate_workflows_in_directory(path):
self.load_workflows_from_file(os.path.join(path, workflow))
def add_child_workflows(self):
for workflow in self.workflows:
playbook_name = workflow.playbook
children = self.workflows[workflow].options.children
for child in children:
workflow_key = _WorkflowKey(
playbook_name,
extract_workflow_name(child, playbook_name=playbook_name))
if workflow_key in self.workflows:
logger.info(
'Adding child workflow {0} to workflow {1}'.format(
child, self.workflows[workflow_key].name))
children[child] = self.workflows[workflow_key]
def add_workflow_scheduled_jobs(self):
"""Schedules the workflow to run based on workflow options.
"""
for workflow in self.workflows:
if (self.workflows[workflow].options.enabled
and self.workflows[workflow].options.scheduler['autorun']
== 'true'):
schedule_type = self.workflows[workflow].options.scheduler[
'type']
schedule = self.workflows[workflow].options.scheduler['args']
self.scheduler.add_job(self.workflows[workflow].execute,
trigger=schedule_type,
replace_existing=True,
**schedule)
logger.info('Added scheduled job for workflow {0}'.format(
self.workflows[workflow].name))
def create_workflow_from_template(self,
playbook_name,
workflow_name,
template_playbook='emptyWorkflow',
template_name='emptyWorkflow'):
"""Creates a workflow from a workflow template.
Args:
playbook_name (str): The name of the new playbook.
workflow_name (str): The name of the new workflow.
template_playbook (str): The name of the playbook template to load. Default is "emptyWorkflow".
template_name (str): The name of the workflow template to load. Default is "emptyWorkflow".
Returns:
True on success, False if otherwise.
"""
path = '{0}{1}{2}.workflow'.format(core.config.paths.templates_path,
sep, template_playbook)
return self.load_workflow_from_file(path=path,
workflow_name=template_name,
name_override=workflow_name,
playbook_override=playbook_name)
def create_playbook_from_template(self,
playbook_name,
template_playbook='emptyWorkflow'):
"""Creates a playbook from a playbook template.
Args:
playbook_name (str): The name of the new playbook.
template_playbook (str): The name of the playbook template to load. Default is "emptyWorkflow".
"""
# TODO: Need a handler for returning workflow key and status
path = '{0}{1}{2}.workflow'.format(core.config.paths.templates_path,
sep, template_playbook)
self.load_workflows_from_file(path=path,
playbook_override=playbook_name)
def remove_workflow(self, playbook_name, workflow_name):
"""Removes a workflow.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow to remove.
Returns:
True on success, False otherwise.
"""
name = _WorkflowKey(playbook_name, workflow_name)
if name in self.workflows:
del self.workflows[name]
logger.debug('Removed workflow {0}'.format(name))
return True
logger.warning(
'Cannot remove workflow {0}. Does not exist in controller'.format(
name))
return False
def remove_playbook(self, playbook_name):
"""Removes a playbook and all workflows within it.
Args:
playbook_name (str): The name of the playbook to remove.
Returns:
True on success, False otherwise.
"""
for name in [
workflow for workflow in self.workflows
if workflow.playbook == playbook_name
]:
del self.workflows[name]
logger.debug('Removed workflow {0}'.format(name))
logger.debug('Removed playbook {0}'.format(playbook_name))
return True
def get_all_workflows(self):
"""Gets all of the currently loaded workflows.
Returns:
A dict with key being the playbook, mapping to a list of workflow names for each playbook.
"""
result = {}
for key in self.workflows.keys():
if key.playbook not in result:
result[key.playbook] = []
result[key.playbook].append(key.workflow)
return result
def get_all_playbooks(self):
"""Gets a list of all playbooks.
Returns:
A list containing all currently loaded playbook names.
"""
return list(set(key.playbook for key in self.workflows.keys()))
def is_workflow_registered(self, playbook_name, workflow_name):
"""Checks whether or not a workflow is currently registered in the system.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow.
Returns:
True if the workflow is registered, false otherwise.
"""
return _WorkflowKey(playbook_name, workflow_name) in self.workflows
def is_playbook_registered(self, playbook_name):
"""Checks whether or not a playbook is currently registered in the system.
Args:
playbook_name (str): The name of the playbook.
Returns:
True if the playbook is registered, false otherwise.
"""
return any(workflow_key.playbook == playbook_name
for workflow_key in self.workflows)
def update_workflow_name(self, old_playbook, old_workflow, new_playbook,
new_workflow):
"""Update the name of a workflow.
Args:
old_playbook (str): Name of the current playbook.
old_workflow (str): Name of the current workflow.
new_playbook (str): The new name of the playbook.
new_workflow (str): The new name of the workflow.
"""
old_key = _WorkflowKey(old_playbook, old_workflow)
new_key = _WorkflowKey(new_playbook, new_workflow)
self.workflows[new_key] = self.workflows.pop(old_key)
self.workflows[new_key].name = construct_workflow_name_key(
new_playbook, new_workflow)
self.workflows[new_key].reconstruct_ancestry([self.name])
logger.debug('updated workflow name {0} to {1}'.format(
old_key, new_key))
def update_playbook_name(self, old_playbook, new_playbook):
"""Update the name of a playbook.
Args:
old_playbook (str): Name of the current playbook.
new_playbook (str): The new name of the playbook.
"""
for key in [
name for name in self.workflows.keys()
if name.playbook == old_playbook
]:
self.update_workflow_name(old_playbook, key.workflow, new_playbook,
key.workflow)
def add_workflow_breakpoint_steps(self, playbook_name, workflow_name,
steps):
"""Adds a breakpoint (for debugging purposes) in the specified steps.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow under which the steps are located.
steps (list[str]): The list of step names for which the user would like to pause execution.
"""
workflow = self.get_workflow(playbook_name, workflow_name)
if workflow:
workflow.breakpoint_steps.extend(steps)
def execute_workflow(self, playbook_name, workflow_name, start=None):
"""Executes a workflow.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): Workflow to execute.
start (str, optional): The name of the first step. Defaults to "start".
"""
global pool
global workflows
global threading_is_initialized
key = _WorkflowKey(playbook_name, workflow_name)
if key in self.workflows:
workflow = self.workflows[key]
subs = deepcopy(subscription.subscriptions)
# If threading has not been initialized, initialize it.
if not threading_is_initialized:
initialize_threading()
if start is not None:
logger.info('Executing workflow {0} for step {1}'.format(
key, start))
workflows.append(
pool.submit(execute_workflow_worker, workflow, subs,
start))
else:
logger.info(
'Executing workflow {0} with default starting step'.format(
key, start))
workflows.append(
pool.submit(execute_workflow_worker, workflow, subs))
callbacks.SchedulerJobExecuted.send(self)
else:
logger.error(
'Attempted to execute playbook which does not exist in controller'
)
def get_workflow(self, playbook_name, workflow_name):
"""Get a workflow object.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow.
Returns:
The workflow object if found, else None.
"""
key = _WorkflowKey(playbook_name, workflow_name)
if key in self.workflows:
return self.workflows[key]
return None
def get_all_workflows_by_playbook(self, playbook_name):
"""Get a list of all workflow objects in a playbook.
Args:
playbook_name: The name of the playbook.
Returns:
A list of all workflow objects in a playbook.
"""
_workflows = []
for key in self.workflows.keys():
if key.playbook == playbook_name:
_workflows.append(self.workflows[key].name)
return _workflows
def playbook_to_xml(self, playbook_name):
"""Returns the XML representation of a playbook.
Args:
playbook_name: The name of the playbook.
Returns:
The XML representation of the playbook if the playbook has any workflows under it, else None.
"""
all_workflows = [
workflow for key, workflow in self.workflows.items()
if key.playbook == playbook_name
]
if all_workflows:
xml = cElementTree.Element('workflows')
for workflow in all_workflows:
xml.append(workflow.to_xml())
return xml
else:
logger.debug(
'No workflows are registered in controller to convert to XML')
return None
def copy_workflow(self, old_playbook_name, new_playbook_name,
old_workflow_name, new_workflow_name):
"""Duplicates a workflow into its current playbook, or a different playbook.
Args:
old_playbook_name (str): Playbook name under which the workflow is located.
new_playbook_name (str): The new playbook name for the duplicated workflow.
old_workflow_name (str): The name of the workflow to be copied.
new_workflow_name (str): The new name of the duplicated workflow.
"""
workflow = self.get_workflow(old_playbook_name, old_workflow_name)
workflow_copy = deepcopy(workflow)
workflow_copy.playbook_name = new_playbook_name
workflow_copy.name = construct_workflow_name_key(
new_playbook_name, new_workflow_name)
key = _WorkflowKey(new_playbook_name, new_workflow_name)
self.workflows[key] = workflow_copy
self.workflows[key].reconstruct_ancestry([self.name])
logger.info('Workflow copied from {0}-{1} to {2}-{3}'.format(
old_playbook_name, old_workflow_name, new_playbook_name,
new_workflow_name))
def copy_playbook(self, old_playbook_name, new_playbook_name):
"""Copies a playbook
Args:
old_playbook_name (str): The name of the playbook to be copied.
new_playbook_name (str): The new name of the duplicated playbook.
"""
for workflow in [
workflow.workflow for workflow in self.workflows
if workflow.playbook == old_playbook_name
]:
self.copy_workflow(old_playbook_name, new_playbook_name, workflow,
workflow)
def pause_workflow(self, playbook_name, workflow_name):
"""Pauses a workflow that is currently executing.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow.
Returns:
A randomly-generated key that needs to be used in order to resume the workflow. This feature is added for
security purposes.
"""
workflow = self.get_workflow(playbook_name, workflow_name)
wf_key = _WorkflowKey(playbook_name, workflow_name)
self.paused_workflows[wf_key] = uuid.uuid4()
if workflow:
logger.info('Pausing workflow {0}'.format(workflow.name))
workflow.pause()
return self.paused_workflows[wf_key].hex
def resume_workflow(self, playbook_name, workflow_name, validate_uuid):
"""Resumes a workflow that has been paused.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow.
validate_uuid (str): The randomly-generated hexadecimal key that was returned from pause_workflow(). This
is needed to resume a workflow for security purposes.
Returns:
"Success" if it is successful, or other error messages.
"""
workflow = self.get_workflow(playbook_name, workflow_name)
wf_key = _WorkflowKey(playbook_name, workflow_name)
if workflow:
if validate_uuid == self.paused_workflows[wf_key].hex:
logger.info('Resuming workflow {0}'.format(workflow.name))
workflow.resume()
return True
else:
logger.warning(
'Cannot resume workflow {0}. Invalid key'.format(
workflow.name))
return False
def resume_breakpoint_step(self, playbook_name, workflow_name):
"""Resumes a step that has been specified as a breakpoint.
Args:
playbook_name (str): Playbook name under which the workflow is located.
workflow_name (str): The name of the workflow.
"""
workflow = self.get_workflow(playbook_name, workflow_name)
if workflow:
logger.debug('Resuming workflow {0} from breakpoint'.format(
workflow.name))
workflow.resume_breakpoint_step()
# Starts active execution
def start(self):
"""Starts the scheduler for active execution. This function must be called before any workflows are executed.
Returns:
The state of the scheduler if successful, error message if scheduler is in "stopped" state.
"""
if self.scheduler.state != STATE_RUNNING and self.scheduler.state != STATE_PAUSED:
logger.info('Starting scheduler')
self.scheduler.start()
else:
logger.warning(
'Cannot start scheduler. Scheduler is already running or is paused'
)
return "Scheduler already running."
return self.scheduler.state
# Stops active execution
def stop(self, wait=True):
"""Stops active execution.
Args:
wait (bool, optional): Boolean to synchronously or asynchronously wait for the scheduler to shutdown.
Default is True.
Returns:
The state of the scheduler if successful, error message if scheduler is already in "stopped" state.
"""
if self.scheduler.state != STATE_STOPPED:
logger.info('Stopping scheduler')
self.scheduler.shutdown(wait=wait)
else:
logger.warning(
'Cannot stop scheduler. Scheduler is already stopped')
return "Scheduler already stopped."
return self.scheduler.state
# Pauses active execution
def pause(self):
"""Pauses active execution.
Returns:
The state of the scheduler if successful, error message if scheduler is not in the "running" state.
"""
if self.scheduler.state == STATE_RUNNING:
logger.info('Pausing scheduler')
self.scheduler.pause()
elif self.scheduler.state == STATE_PAUSED:
logger.warning(
'Cannot pause scheduler. Scheduler is already paused')
return "Scheduler already paused."
elif self.scheduler.state == STATE_STOPPED:
logger.warning('Cannot pause scheduler. Scheduler is stopped')
return "Scheduler is in STOPPED state and cannot be paused."
return self.scheduler.state
# Resumes active execution
def resume(self):
"""Resumes active execution.
Returns:
The state of the scheduler if successful, error message if scheduler is not in the "paused" state.
"""
if self.scheduler.state == STATE_PAUSED:
logger.info('Resuming scheduler')
self.scheduler.resume()
else:
logger.warning(
"Scheduler is not in PAUSED state and cannot be resumed.")
return "Scheduler is not in PAUSED state and cannot be resumed."
return self.scheduler.state
# Pauses active execution of specific job
def pause_job(self, job_id):
"""Pauses active execution of a specific job.
Args:
job_id (str): ID of the job to pause.
"""
logger.info('Pausing job {0}'.format(job_id))
self.scheduler.pause_job(job_id=job_id)
# Resumes active execution of specific job
def resume_job(self, job_id):
"""Resumes active execution of a specific job.
Args:
job_id (str): ID of the job to resume.
"""
logger.info('Resuming job {0}'.format(job_id))
self.scheduler.resume_job(job_id=job_id)
# Returns jobs scheduled for active execution
def get_scheduled_jobs(self):
"""Get all actively scheduled jobs.
Returns:
A list of all actively scheduled jobs.
"""
self.scheduler.get_jobs()
def __scheduler_listener(self):
event_selector_map = {
EVENT_SCHEDULER_START:
(lambda: callbacks.SchedulerStart.send(self)),
EVENT_SCHEDULER_SHUTDOWN:
(lambda: callbacks.SchedulerShutdown.send(self)),
EVENT_SCHEDULER_PAUSED:
(lambda: callbacks.SchedulerPaused.send(self)),
EVENT_SCHEDULER_RESUMED:
(lambda: callbacks.SchedulerResumed.send(self)),
EVENT_JOB_ADDED: (lambda: callbacks.SchedulerJobAdded.send(self)),
EVENT_JOB_REMOVED:
(lambda: callbacks.SchedulerJobRemoved.send(self)),
EVENT_JOB_EXECUTED:
(lambda: callbacks.SchedulerJobExecuted.send(self)),
EVENT_JOB_ERROR: (lambda: callbacks.SchedulerJobError.send(self))
}
def event_selector(event):
try:
event_selector_map[event.code]()
except KeyError:
print('Error: Unknown event sent!')
return event_selector
class Scheduler(object):
def __init__(self):
self.scheduler = GeventScheduler()
self.scheduler.add_listener(self.__scheduler_listener(),
EVENT_SCHEDULER_START | EVENT_SCHEDULER_SHUTDOWN
| EVENT_SCHEDULER_PAUSED | EVENT_SCHEDULER_RESUMED
| EVENT_JOB_ADDED | EVENT_JOB_REMOVED
| EVENT_JOB_EXECUTED | EVENT_JOB_ERROR)
self.id = 'controller'
self.app = None
def schedule_workflows(self, task_id, executable, workflow_ids, trigger):
"""
Schedules a workflow for execution
Args:
task_id (int): Id of the scheduled task
executable (func): A callable to execute must take in one argument -- a workflow id
workflow_ids (iterable(str)): An iterable of workflow ids
trigger (Trigger): The trigger to use for this scheduled task
"""
def execute(id_):
with self.app.app_context():
executable(id_)
for workflow_id in workflow_ids:
self.scheduler.add_job(execute, args=(workflow_id,),
id=construct_task_id(task_id, workflow_id),
trigger=trigger, replace_existing=True)
def get_all_scheduled_workflows(self):
"""
Gets all the scheduled workflows
Returns:
(dict{str: list[str]}) A dict of task_id to workflow execution ids
"""
tasks = {}
for job in self.scheduler.get_jobs():
task, workflow_execution_id = split_task_id(job.id)
if task not in tasks:
tasks[task] = [workflow_execution_id]
else:
tasks[task].append(workflow_execution_id)
return tasks
def get_scheduled_workflows(self, task_id):
"""
Gets all the scheduled worfklows for a given task id
Args:
task_id (str): The task id
Returns:
(list[str]) A list fo workflow execution id associated with this task id
"""
tasks = []
for job in self.scheduler.get_jobs():
task, workflow_execution_id = split_task_id(job.id)
if task == task_id:
tasks.append(workflow_execution_id)
return tasks
def update_workflows(self, task_id, trigger):
"""
Updates the workflows for a given task id to use a different trigger
Args:
task_id (str|int): The task id to update
trigger (Trigger): The new trigger to use
"""
existing_tasks = {construct_task_id(task_id, workflow_execution_id) for workflow_execution_id in
self.get_scheduled_workflows(task_id)}
for job_id in existing_tasks:
self.scheduler.reschedule_job(job_id=job_id, trigger=trigger)
def unschedule_workflows(self, task_id, workflow_execution_ids):
"""
Unschedules a workflow
Args:
task_id (str|int): The task ID to unschedule
workflow_execution_ids (list[str]): The list of workflow execution IDs to update
"""
for workflow_execution_id in workflow_execution_ids:
try:
self.scheduler.remove_job(construct_task_id(task_id, workflow_execution_id))
except JobLookupError:
logger.warning('Cannot delete task {}. '
'No task found in scheduler'.format(construct_task_id(task_id, workflow_execution_id)))
def start(self):
"""Starts the scheduler for active execution. This function must be called before any workflows are executed.
Returns:
The state of the scheduler if successful, error message if scheduler is in "stopped" state.
"""
if self.scheduler.state == STATE_STOPPED:
logger.info('Starting scheduler')
self.scheduler.start()
else:
logger.warning('Cannot start scheduler. Scheduler is already running or is paused')
return "Scheduler already running."
return self.scheduler.state
def stop(self, wait=True):
"""Stops active execution.
Args:
wait (bool, optional): Boolean to synchronously or asynchronously wait for the scheduler to shutdown.
Default is True.
Returns:
The state of the scheduler if successful, error message if scheduler is already in "stopped" state.
"""
if self.scheduler.state != STATE_STOPPED:
logger.info('Stopping scheduler')
self.scheduler.shutdown(wait=wait)
else:
logger.warning('Cannot stop scheduler. Scheduler is already stopped')
return "Scheduler already stopped."
return self.scheduler.state
def pause(self):
"""Pauses active execution.
Returns:
The state of the scheduler if successful, error message if scheduler is not in the "running" state.
"""
if self.scheduler.state == STATE_RUNNING:
logger.info('Pausing scheduler')
self.scheduler.pause()
elif self.scheduler.state == STATE_PAUSED:
logger.warning('Cannot pause scheduler. Scheduler is already paused')
return "Scheduler already paused."
elif self.scheduler.state == STATE_STOPPED:
logger.warning('Cannot pause scheduler. Scheduler is stopped')
return "Scheduler is in STOPPED state and cannot be paused."
return self.scheduler.state
def resume(self):
"""Resumes active execution.
Returns:
The state of the scheduler if successful, error message if scheduler is not in the "paused" state.
"""
if self.scheduler.state == STATE_PAUSED:
logger.info('Resuming scheduler')
self.scheduler.resume()
else:
logger.warning("Scheduler is not in PAUSED state and cannot be resumed.")
return "Scheduler is not in PAUSED state and cannot be resumed."
return self.scheduler.state
def pause_workflows(self, task_id, workflow_execution_ids):
"""
Pauses some workflows associated with a task
Args:
task_id (int|str): The id of the task to pause
workflow_execution_ids (list[str]): The list of workflow execution IDs to pause
"""
for workflow_execution_id in workflow_execution_ids:
job_id = construct_task_id(task_id, workflow_execution_id)
try:
self.scheduler.pause_job(job_id=job_id)
logger.info('Paused job {0}'.format(job_id))
except JobLookupError:
logger.warning('Cannot pause scheduled workflow {}. Workflow ID not found'.format(job_id))
def resume_workflows(self, task_id, workflow_execution_ids):
"""
Resumes some workflows associated with a task
Args:
task_id (int|str): The id of the task to pause
workflow_execution_ids (list[str]): The list of workflow execution IDs to resume
"""
for workflow_execution_id in workflow_execution_ids:
job_id = construct_task_id(task_id, workflow_execution_id)
try:
self.scheduler.resume_job(job_id=job_id)
logger.info('Resumed job {0}'.format(job_id))
except JobLookupError:
logger.warning('Cannot resume scheduled workflow {}. Workflow ID not found'.format(job_id))
def __scheduler_listener(self):
event_selector_map = {EVENT_SCHEDULER_START: WalkoffEvent.SchedulerStart,
EVENT_SCHEDULER_SHUTDOWN: WalkoffEvent.SchedulerShutdown,
EVENT_SCHEDULER_PAUSED: WalkoffEvent.SchedulerPaused,
EVENT_SCHEDULER_RESUMED: WalkoffEvent.SchedulerResumed,
EVENT_JOB_ADDED: WalkoffEvent.SchedulerJobAdded,
EVENT_JOB_REMOVED: WalkoffEvent.SchedulerJobRemoved,
EVENT_JOB_EXECUTED: WalkoffEvent.SchedulerJobExecuted,
EVENT_JOB_ERROR: WalkoffEvent.SchedulerJobError}
def event_selector(event):
try:
event = event_selector_map[event.code]
event.send(self)
except KeyError: # pragma: no cover
logger.error('Unknown event sent triggered in scheduler {}'.format(event))
return event_selector