├── .gitignore ├── LICENSE ├── README.md ├── example ├── setup.py └── taskflow_basic_example │ ├── main.py │ ├── tasks │ ├── __init__.py │ └── vacuum_postgres.py │ └── workflows │ ├── __init__.py │ └── echo_workflow.py ├── setup.py ├── taskflow ├── __init__.py ├── alembic.ini ├── cli.py ├── core │ ├── __init__.py │ ├── models.py │ ├── pusher.py │ ├── scheduler.py │ └── worker.py ├── db.py ├── migrations │ ├── README │ ├── env.py │ ├── script.py.mako │ └── versions │ │ ├── 1e64f390802b_add_unique_task_and_workflow_indexes.py │ │ ├── 6ef79b56ad4a_rename_status_retry.py │ │ └── 94e7b91d83d5_create_task_workflow_name_column.py ├── monitoring │ ├── __init__.py │ ├── aws.py │ ├── base.py │ └── slack.py ├── push_workers │ ├── __init__.py │ ├── aws_batch.py │ └── base.py ├── rest │ ├── __init__.py │ ├── app.py │ └── resources.py └── tasks │ ├── __init__.py │ └── bash_task.py └── tests ├── shared_fixtures.py ├── test_basic_workflow.py ├── test_monitoring.py ├── test_push_aws_batch.py ├── test_rest.py └── test_tasks.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | *.egg-info/ 23 | .installed.cfg 24 | *.egg 25 | 26 | # PyInstaller 27 | # Usually these files are written by a python script from a template 28 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 29 | *.manifest 30 | *.spec 31 | 32 | # Installer logs 33 | pip-log.txt 34 | pip-delete-this-directory.txt 35 | 36 | # Unit test / coverage reports 37 | htmlcov/ 38 | .tox/ 39 | .coverage 40 | .coverage.* 41 | .cache 42 | nosetests.xml 43 | coverage.xml 44 | *,cover 45 | 46 | # Translations 47 | *.mo 48 | *.pot 49 | 50 | # Django stuff: 51 | *.log 52 | 53 | # Sphinx documentation 54 | docs/_build/ 55 | 56 | # PyBuilder 57 | target/ 58 | 59 | # pyenv 60 | .python-version 61 | 62 | # dotenv 63 | .env 64 | 65 | node_modules/ 66 | *.log 67 | 68 | # Mac 69 | .DS_Store 70 | ._* 71 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2017 City Of Philadelphia 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Taskflow 2 | 3 | An advanced yet simple system to run your background tasks and workflows. 4 | 5 | Features 6 | 7 | - Recurring tasks (aka jobs) and workflows (a series of dependent tasks) with CRON-like scheduling 8 | 9 | 10 | - Workflow dependencies - tasks execute in order and/or in parallel depending on dependency chain. 11 | 12 | - Two types of workers that execute tasks 13 | - Pull workers - Pull tasks directly off the database queue and execute them. 14 | - Push workers - Tasks are pushed to a remote work management system, such as AWS Batch, Kubernetes, or Apache Mesos. Container friendly. 15 | 16 | ## Motivation 17 | 18 | Other background task and workflow management solutions exist out there, such as Celery for tasks/jobs or Airflow and Luigi for workflows. Taskflow is designed to have a small footprint, maintain it's state in a readily queryable SQL database, have predictable CRON-like scheduling behavior, and GTFO of the way when you need it to. 19 | 20 | ## Concepts 21 | 22 | ### Task 23 | 24 | A Task represents some runnable code or command. For example, extact a table from a database, or push to an API. Anything that can run in pyhon or be excuted in Bash. 25 | 26 | ### Task Instances 27 | 28 | A Task Instance is a specific run of a Task. Task instances can be created programmatically on-demand or automatically using a recurring schedule attached to a Task. 29 | 30 | #### Workflows 31 | 32 | A Workflow represents a series of dependent tasks represent in a graph. Workflows are associated to the Tasks they run. 33 | 34 | #### Workflow Instances 35 | 36 | A Workflow Instance is a specific run of a Workflow. A Workflow Instance creates Task Instances as needed during a run. Like a Task Instance, Workflow Instances can be created programmatically on-demand or automatically using a recurring schedule attached to a Workflow. 37 | 38 | #### Scheduler 39 | 40 | The scheduler always runs as a single instance at a time. It schedules recurring Tasks and recurring Workflows. It also advances running Workflow Instances, scheduling Task Instances as needed. 41 | 42 | #### Pusher 43 | 44 | The Pusher is usually run within the same process as the scheduler. The Pusher pulls tasks destined for a push worker off the task_instances table and pushes them to the push destination. For examples, pushing tasks to AWS Batch. The Push also syncs the state of the currently pushed tasks with the push destination. Multiple push destinations can be used at the same time, for example one task could go to AWS Batch while another goes to Kubernetes. 45 | 46 | #### Pull Worker 47 | 48 | A pull worker is a process that directly pulls tasks off the queue and executes them. 49 | 50 | ## API 51 | 52 | ### Task 53 | 54 | Creates a new Task definiton that can be used to schedule tasks and create task instances. 55 | 56 | ```python 57 | task_foo = Task( 58 | name=None, # Required - name of the task, like 'push_new_accounts' 59 | workflow=None, # workflow the task is associated with 60 | active=False, # wether the task is active, if false, it will not be sceduled 61 | title=None, # human friendly title of the task, used for the UI 62 | description=None, # human friendly description of the task, used for the UI 63 | concurrency=1, # maximum number of TaskInstances of this Task that can be running at a time 64 | sla=None, # numbers of seconds the scheduled Task should start within 65 | schedule=None, # CRON pattern to schedule the Task, if the task is to be scheduled 66 | default_priority='normal', # default priority used for the TaskInstances of this Task 67 | start_date=None, # start datetime for scheduling this Task, the task will not be scheduled before this time 68 | end_date=None, # end datetime for scheduling this Task, the task will not be scheduled after this time 69 | retries=0, # number of times to retry task if it fails 70 | timeout=300, # if the task is not completed and the locked_at datetime is past these number of seconds, retry or fail task 71 | retry_delay=300, # number of seconds inbetween retries 72 | params={}, # parameters of the task, these are available to the executors and pushers 73 | push_destination=None) # for Tasks that need to be pushed, where to push them, for example 'aws_batch' 74 | 75 | task_foo.depends_on(task_bar) # Task.depends_on tells Taskflow that one Task is dependant on another. Both tasks need to be in the same Workflow 76 | ``` 77 | 78 | To create a task with executable python code, override the `execute` function on an inherited class. 79 | 80 | ```python 81 | class MyTask(Task): 82 | def execute(self, task_instance): 83 | print('Any sort of python I want') 84 | ``` 85 | 86 | ### Workflow 87 | 88 | Creates a new Workflow definiton that can be used to schedule workflows and create workflow instances. 89 | 90 | ```python 91 | workflow_foo = Workflow( 92 | name=None, # Required. Name of the workflow like 'salesforce_etl' 93 | active=False, # wether the workflow is active, if false, it will not be sceduled 94 | title=None, # human friendly title of the workflow, used for the UI 95 | description=None, # human friendly description of the workflow, used for the UI 96 | concurrency=1, # maximum number of WorkflowInstances of this Workflow that can be running at a time 97 | sla=None, # numbers of seconds the scheduled Workflow should start within 98 | schedule=None, # CRON pattern to schedule the Workflow, if the workflow is to be scheduled 99 | default_priority='normal', # default priority used for the WorkflowInstances of this Workflow 100 | start_date=None, # start datetime for scheduling this Workflow, the workflow will not be scheduled before this time 101 | end_date=None, # end datetime for scheduling this Workflow, the workflow will not be scheduled after this time) 102 | ``` 103 | 104 | ### Taskflow 105 | 106 | The Taskflow class is used to create an instance of Taskflow and associate your Tasks and Workflows. 107 | 108 | ```python 109 | taskflow = Taskflow() 110 | 111 | # These function take arrays of Workflows and Tasks (not associated to a Workflow) and add them to Taskflow 112 | taskflow.add_workflows(workflows) 113 | taskflow.add_tasks(tasks) 114 | ``` 115 | 116 | The above code is usually all you need to setup Taskflow, see the exmaple directory for an example implementation. 117 | 118 | ## CLI 119 | 120 | The CLI is used to manage Taskflow, manage the database, run the scheduler, run the pull worker, and execute task instances. 121 | 122 | ``` 123 | Usage: main.py [OPTIONS] COMMAND [ARGS]... 124 | 125 | Options: 126 | --taskflow TEXT 127 | --help Show this message and exit. 128 | 129 | Commands: 130 | api_server 131 | init_db 132 | migrate_db 133 | pull_worker 134 | queue_task 135 | queue_workflow 136 | run_task 137 | scheduler 138 | ``` 139 | -------------------------------------------------------------------------------- /example/setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from distutils.core import setup 4 | 5 | setup( 6 | name='taskflow_basic_example', 7 | version='0.1dev', 8 | packages=[ 9 | 'taskflow_basic_example', 10 | 'taskflow_basic_example.workflows', 11 | 'taskflow_basic_example.tasks' 12 | ], 13 | install_requires=[ 14 | ], 15 | dependency_links=[ 16 | ], 17 | ) 18 | -------------------------------------------------------------------------------- /example/taskflow_basic_example/main.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import logging 4 | from datetime import datetime 5 | import json 6 | 7 | from taskflow import cli, Taskflow 8 | 9 | from taskflow_basic_example.workflows import workflows 10 | from taskflow_basic_example.tasks import tasks 11 | 12 | DEBUG = os.getenv('DEBUG', 'False') == 'True' 13 | 14 | def get_logging(): 15 | logger = logging.getLogger() 16 | handler = logging.StreamHandler() 17 | formatter = logging.Formatter('[%(asctime)s] %(name)s %(levelname)s %(message)s') 18 | handler.setFormatter(formatter) 19 | logger.addHandler(handler) 20 | 21 | if DEBUG: 22 | level = logging.DEBUG 23 | else: 24 | level = logging.INFO 25 | 26 | logger.setLevel(level) 27 | 28 | def get_taskflow(): 29 | taskflow = Taskflow() 30 | 31 | taskflow.add_workflows(workflows) 32 | taskflow.add_tasks(tasks) 33 | 34 | return taskflow 35 | 36 | if __name__ == '__main__': 37 | get_logging() 38 | taskflow = get_taskflow() 39 | cli(taskflow) 40 | -------------------------------------------------------------------------------- /example/taskflow_basic_example/tasks/__init__.py: -------------------------------------------------------------------------------- 1 | from .vacuum_postgres import vacuum_events, vacuum_aggregates 2 | 3 | tasks = [ 4 | vacuum_events, 5 | vacuum_aggregates 6 | ] 7 | -------------------------------------------------------------------------------- /example/taskflow_basic_example/tasks/vacuum_postgres.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from taskflow.tasks.bash_task import BashTask 4 | 5 | class VacuumPostgres(BashTask): 6 | def get_command(self): 7 | return 'psql -c "VACUUM ANALYZE {}" -h {} -p {} -u "$POSTGRES_USERNAME"'.format( 8 | self.params['table_name'], 9 | self.params['host'], 10 | self.params['port']) 11 | 12 | vacuum_events = VacuumPostgres( 13 | name='vacuum_events', 14 | active=True, 15 | schedule='0 8 * * 2-6', # Every Tuesday through Saturday at 8am UTC 16 | timeout=3600, 17 | params={ 18 | 'table_name': 'events', 19 | 'host': 'localhost', 20 | 'port': 5432 21 | }) 22 | 23 | vacuum_aggregates = VacuumPostgres( 24 | name='vacuum_aggregates', 25 | active=True, 26 | schedule='0 8 * * 2-6', # Every Tuesday through Saturday at 8am UTC 27 | timeout=3600, 28 | params={ 29 | 'table_name': 'user_aggregate_metrics', 30 | 'host': 'localhost', 31 | 'port': 5432 32 | }) 33 | -------------------------------------------------------------------------------- /example/taskflow_basic_example/workflows/__init__.py: -------------------------------------------------------------------------------- 1 | from .echo_workflow import echo_workflow 2 | 3 | workflows = [ 4 | echo_workflow 5 | ] 6 | -------------------------------------------------------------------------------- /example/taskflow_basic_example/workflows/echo_workflow.py: -------------------------------------------------------------------------------- 1 | from taskflow import Workflow 2 | from taskflow.tasks.bash_task import BashTask 3 | 4 | echo_workflow = Workflow( 5 | name='echo_workflow', 6 | active=True, 7 | schedule='0 * * * *') 8 | 9 | echo_a = BashTask(workflow=echo_workflow, 10 | name='echo_a', 11 | active=True, 12 | retries=2, # this task is allowed to retry 13 | params={ 14 | 'command': 'echo "Foo A"' 15 | }) 16 | 17 | echo_b = BashTask(workflow=echo_workflow, 18 | name='echo_b', 19 | active=True, 20 | retries=2, # this task is allowed to retry 21 | retry_delay=1200, # this task can be expensive, lets wait 20 minutes between retries 22 | params={ 23 | 'command': 'echo "Foo B"' 24 | }) 25 | 26 | echo_c = BashTask(workflow=echo_workflow, 27 | name='echo_c', 28 | active=True, 29 | params={ 30 | 'command': 'echo "Foo C"' 31 | }) 32 | 33 | echo_d = BashTask(workflow=echo_workflow, 34 | name='echo_d', 35 | active=True, 36 | params={ 37 | 'command': 'echo "Foo D"' 38 | }) 39 | 40 | echo_c.depends_on(echo_a) 41 | echo_c.depends_on(echo_b) 42 | echo_d.depends_on(echo_c) 43 | 44 | ## A and B should run in parallel, while C waits, D should then wait for C 45 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from distutils.core import setup 4 | from setuptools import find_packages 5 | 6 | setup( 7 | name='taskflow', 8 | version='0.2.13', 9 | packages=find_packages(), 10 | install_requires=[ 11 | 'alembic==0.9.6', 12 | 'boto3==1.4.4', 13 | 'click==6.7', 14 | 'cron-descriptor==1.2.10', 15 | 'croniter==0.3.17', 16 | 'Flask==0.12.2', 17 | 'Flask-Cors==3.0.2', 18 | 'Flask-RESTful==0.3.6', 19 | 'Flask-SQLAlchemy==2.2', 20 | 'gunicorn==19.7.1', 21 | 'marshmallow==2.13.5', 22 | 'psycopg2==2.7.1', 23 | 'pytest==3.1.1', 24 | 'restful_ben==0.1.0', 25 | 'requests==2.17.3', 26 | 'smart-open==1.5.5', 27 | 'SQLAlchemy==1.1.10', 28 | 'toposort==1.5' 29 | ], 30 | dependency_links=[ 31 | 'https://github.com/CityOfPhiladelphia/restful-ben/tarball/0.0.1#egg=restful_ben-0.1.0' 32 | ], 33 | ) 34 | -------------------------------------------------------------------------------- /taskflow/__init__.py: -------------------------------------------------------------------------------- 1 | from .core.models import Taskflow, Task, Workflow, WorkflowInstance, TaskInstance 2 | from .core.scheduler import Scheduler 3 | from .core.pusher import Pusher 4 | from .core.worker import Worker 5 | from .cli import cli 6 | -------------------------------------------------------------------------------- /taskflow/alembic.ini: -------------------------------------------------------------------------------- 1 | # A generic, single database configuration. 2 | 3 | [alembic] 4 | # path to migration scripts 5 | script_location = migrations 6 | 7 | # template used to generate migration files 8 | # file_template = %%(rev)s_%%(slug)s 9 | 10 | # timezone to use when rendering the date 11 | # within the migration file as well as the filename. 12 | # string value is passed to dateutil.tz.gettz() 13 | # leave blank for localtime 14 | timezone = UTC 15 | 16 | # max length of characters to apply to the 17 | # "slug" field 18 | #truncate_slug_length = 40 19 | 20 | # set to 'true' to run the environment during 21 | # the 'revision' command, regardless of autogenerate 22 | # revision_environment = false 23 | 24 | # set to 'true' to allow .pyc and .pyo files without 25 | # a source .py file to be detected as revisions in the 26 | # versions/ directory 27 | # sourceless = false 28 | 29 | # version location specification; this defaults 30 | # to migrations/versions. When using multiple version 31 | # directories, initial revisions must be specified with --version-path 32 | # version_locations = %(here)s/bar %(here)s/bat migrations/versions 33 | 34 | # the output encoding used when revision files 35 | # are written from script.py.mako 36 | # output_encoding = utf-8 37 | 38 | sqlalchemy.url = driver://user:pass@localhost/dbname 39 | 40 | 41 | # Logging configuration 42 | [loggers] 43 | keys = root,sqlalchemy,alembic 44 | 45 | [handlers] 46 | keys = console 47 | 48 | [formatters] 49 | keys = generic 50 | 51 | [logger_root] 52 | level = WARN 53 | handlers = console 54 | qualname = 55 | 56 | [logger_sqlalchemy] 57 | level = WARN 58 | handlers = 59 | qualname = sqlalchemy.engine 60 | 61 | [logger_alembic] 62 | level = INFO 63 | handlers = 64 | qualname = alembic 65 | 66 | [handler_console] 67 | class = StreamHandler 68 | args = (sys.stderr,) 69 | level = NOTSET 70 | formatter = generic 71 | 72 | [formatter_generic] 73 | format = %(levelname)-5.5s [%(name)s] %(message)s 74 | datefmt = %H:%M:%S 75 | -------------------------------------------------------------------------------- /taskflow/cli.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import logging 4 | from datetime import datetime 5 | import json 6 | import os 7 | import socket 8 | import sys 9 | import multiprocessing 10 | 11 | import requests 12 | import click 13 | from sqlalchemy import create_engine 14 | from sqlalchemy.orm import sessionmaker 15 | import gunicorn.app.base 16 | from gunicorn.six import iteritems 17 | 18 | from taskflow import Scheduler, Pusher, Taskflow, Worker, TaskInstance 19 | from taskflow import db 20 | from taskflow.rest.app import create_app 21 | 22 | def get_logging(): 23 | logger = logging.getLogger() 24 | handler = logging.StreamHandler() 25 | formatter = logging.Formatter('[%(asctime)s] %(name)s %(levelname)s %(message)s') 26 | handler.setFormatter(formatter) 27 | logger.addHandler(handler) 28 | logger.setLevel(logging.DEBUG) 29 | 30 | def exception_handler(type, value, tb): 31 | logger.exception("Uncaught exception: {0}".format(str(value))) 32 | 33 | sys.excepthook = exception_handler 34 | 35 | def get_worker_id(): 36 | worker_components = [] 37 | 38 | ## AWS 39 | try: 40 | response = requests.get('http://169.254.169.254/latest/meta-data/instance-id', timeout=0.1) 41 | if response.status_code == 200: 42 | worker_components.append(response.text) 43 | except: 44 | pass 45 | 46 | ## ECS (AWS Batch uses ECS as well) 47 | try: 48 | response = requests.get('http://172.17.0.1:51678/v1/tasks', timeout=0.1) 49 | if response.status_code == 200: 50 | tasks = response.json()['Tasks'] 51 | short_docker_id = os.getenv('HOSTNAME', None) ## ECS marks the short docker id as the HOSTNAME 52 | if short_docker_id != None: 53 | matched = list(filter( 54 | lambda ecs_task: ecs_task['Containers'][0]['DockerId'][0:12] == short_docker_id, 55 | tasks)) 56 | if len(matched) > 0: 57 | worker_components.append(matched[0]['Containers'][0]['Arn']) 58 | except: 59 | pass 60 | 61 | ## fallback to IP 62 | if len(worker_components) == 0: 63 | return socket.gethostbyname(socket.gethostname()) 64 | else: 65 | return '-'.join(worker_components) 66 | 67 | def number_of_workers(): 68 | return (multiprocessing.cpu_count() * 2) + 1 69 | 70 | class StandaloneApplication(gunicorn.app.base.BaseApplication): 71 | 72 | def __init__(self, app, options=None): 73 | self.options = options or {} 74 | self.application = app 75 | super(StandaloneApplication, self).__init__() 76 | 77 | def load_config(self): 78 | config = dict([(key, value) for key, value in iteritems(self.options) 79 | if key in self.cfg.settings and value is not None]) 80 | for key, value in iteritems(config): 81 | self.cfg.set(key.lower(), value) 82 | 83 | def load(self): 84 | return self.application 85 | 86 | @click.group() 87 | @click.option('--taskflow') 88 | @click.pass_context 89 | def main(ctx, taskflow): 90 | if taskflow != None: 91 | ## TODO: load taskflow from option using dynamic import 92 | ctx.obj['taskflow'] = taskflow 93 | 94 | @main.command() 95 | @click.option('--sql-alchemy-connection') 96 | @click.option('--bind-host', default='0.0.0.0') 97 | @click.option('--bind-port', default='5000', type=int) 98 | @click.option('--worker-class', default='sync') 99 | @click.option('--prod', is_flag=True, default=False) 100 | @click.pass_context 101 | def api_server(ctx, sql_alchemy_connection, bind_host, bind_port, worker_class, prod): 102 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 103 | 104 | taskflow = ctx.obj['taskflow'] 105 | 106 | engine = create_engine(connection_string) 107 | Session = sessionmaker(bind=engine) 108 | session = Session() 109 | taskflow.sync_db(session, read_only=True) 110 | session.close() 111 | 112 | app = create_app(taskflow, connection_string=connection_string) 113 | 114 | if prod: 115 | options = { 116 | 'bind': '{}:{}'.format(bind_host, bind_port), 117 | 'workers': number_of_workers(), 118 | 'worker_class': worker_class 119 | } 120 | StandaloneApplication(app, options).run() 121 | else: 122 | app.run(host=bind_host, port=bind_port) 123 | 124 | @main.command() 125 | @click.option('--sql-alchemy-connection') 126 | @click.pass_context 127 | def init_db(ctx, sql_alchemy_connection): 128 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 129 | 130 | db.init_db(connection_string) 131 | 132 | @main.command() 133 | @click.option('--sql-alchemy-connection') 134 | @click.pass_context 135 | def migrate_db(ctx, sql_alchemy_connection): 136 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 137 | 138 | db.migrate_db(connection_string) 139 | 140 | @main.command() 141 | @click.option('--sql-alchemy-connection') 142 | @click.option('-n','--num-runs', type=int, default=10) 143 | @click.option('--dry-run', is_flag=True, default=False) 144 | @click.option('--now-override') 145 | @click.option('--sleep', type=int, default=5) 146 | @click.pass_context 147 | def scheduler(ctx, sql_alchemy_connection, num_runs, dry_run, now_override, sleep): 148 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 149 | engine = create_engine(connection_string) 150 | Session = sessionmaker(bind=engine) 151 | 152 | session = Session() 153 | taskflow = ctx.obj['taskflow'] 154 | taskflow.sync_db(session) 155 | session.close() 156 | 157 | if now_override != None: 158 | now_override = datetime.strptime(now_override, '%Y-%m-%dT%H:%M:%S') 159 | 160 | scheduler = Scheduler(taskflow, dry_run=dry_run, now_override=now_override) 161 | pusher = Pusher(taskflow, dry_run=dry_run, now_override=now_override) 162 | 163 | ## TODO: fix interrupt 164 | 165 | for n in range(0, num_runs): 166 | if n > 0 and sleep > 0: 167 | time.sleep(sleep) 168 | 169 | session = Session() 170 | taskflow.sync_db(session) 171 | scheduler.run(session) 172 | session.close() 173 | 174 | session = Session() 175 | taskflow.sync_db(session) 176 | pusher.run(session) 177 | session.close() 178 | 179 | @main.command() 180 | @click.option('--sql-alchemy-connection') 181 | @click.option('-n','--num-runs', type=int, default=10) 182 | @click.option('--dry-run', is_flag=True, default=False) 183 | @click.option('--now-override') 184 | @click.option('--sleep', type=int, default=5) 185 | @click.option('--task-names') 186 | @click.option('--worker-id') 187 | @click.pass_context 188 | def pull_worker(ctx, sql_alchemy_connection, num_runs, dry_run, now_override, sleep, task_names, worker_id): 189 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 190 | engine = create_engine(connection_string) 191 | Session = sessionmaker(bind=engine) 192 | 193 | session = Session() 194 | taskflow = ctx.obj['taskflow'] 195 | taskflow.sync_db(session) 196 | session.close() 197 | 198 | worker = Worker(taskflow) 199 | 200 | if now_override != None: 201 | now_override = datetime.strptime(now_override, '%Y-%m-%dT%H:%M:%S') 202 | 203 | if task_names != None: 204 | task_names = task_names.split(',') 205 | 206 | if worker_id == None: 207 | worker_id = get_worker_id() 208 | 209 | for n in range(0, num_runs): 210 | if n > 0 and sleep > 0: 211 | time.sleep(sleep) 212 | 213 | session = Session() 214 | 215 | task_instances = taskflow.pull(session, worker_id, task_names=task_names, now=now_override) 216 | 217 | if len(task_instances) > 0: 218 | worker.execute(session, task_instances[0]) 219 | 220 | session.close() 221 | 222 | @main.command() 223 | @click.argument('task_instance_id', type=int) 224 | @click.option('--sql-alchemy-connection') 225 | @click.option('--worker-id') 226 | @click.pass_context 227 | def run_task(ctx, task_instance_id, sql_alchemy_connection, worker_id): 228 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 229 | engine = create_engine(connection_string) 230 | Session = sessionmaker(bind=engine) 231 | 232 | if worker_id == None: 233 | worker_id = get_worker_id() 234 | 235 | session = Session() 236 | 237 | taskflow = ctx.obj['taskflow'] 238 | taskflow.sync_db(session) 239 | 240 | task_instance = session.query(TaskInstance).get(task_instance_id) 241 | 242 | task_instance.worker_id = worker_id 243 | task_instance.locked_at = datetime.utcnow() 244 | session.commit() 245 | 246 | worker = Worker(taskflow) 247 | success = worker.execute(session, task_instance) 248 | 249 | session.close() 250 | 251 | if not success: 252 | sys.exit(1) 253 | 254 | @main.command() 255 | @click.argument('task_name') 256 | @click.option('--workflow-instance-id') 257 | @click.option('--run-at') 258 | @click.option('--priority') 259 | @click.option('--params') 260 | @click.option('--sql-alchemy-connection') 261 | @click.pass_context 262 | def queue_task(ctx, task_name, workflow_instance_id, run_at, priority, params, sql_alchemy_connection): 263 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 264 | engine = create_engine(connection_string) 265 | Session = sessionmaker(bind=engine) 266 | 267 | session = Session() 268 | 269 | taskflow = ctx.obj['taskflow'] 270 | taskflow.sync_db(session) 271 | 272 | task = taskflow.get_task(task_name) 273 | 274 | if task == None: 275 | raise Exception('Task `{}` not found'.format(task_name)) 276 | 277 | if params != None: 278 | params = json.loads(params) 279 | 280 | task_instance = task.get_new_instance( 281 | run_at=run_at, 282 | workflow_instance_id=workflow_instance_id, 283 | priority=priority, 284 | params=params) 285 | 286 | session.add(task_instance) 287 | session.commit() 288 | session.close() 289 | 290 | @main.command() 291 | @click.argument('workflow_name') 292 | @click.option('--run-at') 293 | @click.option('--priority') 294 | @click.option('--sql-alchemy-connection') 295 | @click.pass_context 296 | def queue_workflow(ctx, workflow_name, run_at, priority, sql_alchemy_connection): 297 | connection_string = sql_alchemy_connection or os.getenv('SQL_ALCHEMY_CONNECTION') 298 | engine = create_engine(connection_string) 299 | Session = sessionmaker(bind=engine) 300 | 301 | session = Session() 302 | 303 | taskflow = ctx.obj['taskflow'] 304 | taskflow.sync_db(session) 305 | 306 | workflow = taskflow.get_workflow(workflow_name) 307 | 308 | if workflow == None: 309 | raise Exception('Workflow `{}` not found'.format(workflow_name)) 310 | 311 | workflow_instance = workflow.get_new_instance( 312 | run_at=run_at, 313 | priority=priority) 314 | 315 | session.add(workflow_instance) 316 | session.commit() 317 | session.close() 318 | 319 | def cli(taskflow): 320 | return main(obj={'taskflow': taskflow}) 321 | 322 | if __name__ == '__main__': 323 | get_logging() 324 | main(obj={}) 325 | -------------------------------------------------------------------------------- /taskflow/core/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CityOfPhiladelphia/taskflow/5c86b50c6d8bce43a04fba1b856b1d56f3c36b6a/taskflow/core/__init__.py -------------------------------------------------------------------------------- /taskflow/core/models.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | from sqlalchemy import ( 4 | Column, 5 | Integer, 6 | BigInteger, 7 | String, 8 | DateTime, 9 | Boolean, 10 | Enum, 11 | Index, 12 | func, 13 | text, 14 | ForeignKey, 15 | MetaData 16 | ) 17 | from sqlalchemy.orm import relationship 18 | from sqlalchemy.event import listens_for 19 | from sqlalchemy.dialects.postgresql import JSONB 20 | from sqlalchemy.ext.declarative import declarative_base, declared_attr 21 | from croniter import croniter 22 | from restful_ben.auth import UserAuthMixin 23 | from flask_login import UserMixin 24 | 25 | from taskflow.monitoring.base import Monitor 26 | 27 | metadata = MetaData() 28 | BaseModel = declarative_base(metadata=metadata) 29 | 30 | class Schedulable(object): 31 | name = Column(String, primary_key=True) 32 | active = Column(Boolean, nullable=False) 33 | 34 | def __init__( 35 | self, 36 | name=None, 37 | active=False, 38 | title=None, 39 | description=None, 40 | concurrency=1, 41 | sla=None, 42 | schedule=None, 43 | default_priority='normal', 44 | start_date=None, 45 | end_date=None): 46 | 47 | self.name = name 48 | if not self.name: 49 | raise Exception('`name` required for {}'.format(self.__class__.__name__)) 50 | 51 | self.active = active 52 | self.title = title 53 | self.description = description 54 | self.concurrency = concurrency 55 | self.sla = sla 56 | 57 | self.schedule = schedule 58 | self.default_priority = default_priority 59 | self.start_date = start_date 60 | self.end_date = end_date 61 | 62 | def next_run(self, base_time=None): 63 | if not base_time: 64 | base_time = datetime.utcnow() 65 | iter = croniter(self.schedule, base_time) 66 | return iter.get_next(datetime) 67 | 68 | def last_run(self, base_time=None): 69 | if not base_time: 70 | base_time = datetime.utcnow() 71 | iter = croniter(self.schedule, base_time) 72 | return iter.get_prev(datetime) 73 | 74 | class Workflow(Schedulable, BaseModel): 75 | __tablename__ = 'workflows' 76 | 77 | _tasks = None 78 | 79 | def __init__(self, *args, **kwargs): 80 | super(Workflow, self).__init__(*args, **kwargs) 81 | 82 | self._tasks = set() 83 | 84 | def __repr__(self): 85 | return ''.format(self.name, self.active) 86 | 87 | ## TODO: remove deactivated tasks from graph ? 88 | def get_dependencies_graph(self): 89 | graph = dict() 90 | for task in self._tasks: 91 | graph[task.name] = task._dependencies 92 | return graph 93 | 94 | def get_tasks(self): 95 | return self._tasks 96 | 97 | def get_task(self, task_name): 98 | for task in self._tasks: 99 | if task.name == task_name: 100 | return task 101 | 102 | def get_new_instance(self, scheduled=False, status='queued', run_at=None, priority=None, unique=None): 103 | return WorkflowInstance( 104 | workflow_name=self.name, 105 | scheduled=scheduled, 106 | run_at=run_at or datetime.utcnow(), 107 | status=status, 108 | priority=priority or self.default_priority, 109 | unique=unique) 110 | 111 | class Task(Schedulable, BaseModel): 112 | __tablename__ = 'tasks' 113 | 114 | workflow_name = Column(String, ForeignKey('workflows.name')) 115 | 116 | def __init__( 117 | self, 118 | workflow=None, 119 | retries=0, 120 | timeout=300, 121 | retry_delay=300, 122 | params={}, 123 | push_destination=None, 124 | *args, **kwargs): 125 | super(Task, self).__init__(*args, **kwargs) 126 | 127 | self.workflow = workflow 128 | if self.workflow: 129 | if self in self.workflow._tasks: 130 | raise Exception('`{}` already added to workflow `{}`'.format(self.name, self.workflow.name)) 131 | self.workflow._tasks.add(self) 132 | self.workflow_name = workflow.name 133 | 134 | self.retries = retries 135 | self.timeout = timeout 136 | self.retry_delay = retry_delay 137 | 138 | self.params = params 139 | 140 | self.push_destination = push_destination 141 | 142 | self._dependencies = set() 143 | 144 | def __repr__(self): 145 | return ''.format(self.name, self.active) 146 | 147 | def depends_on(self, task): 148 | if self.workflow == None: 149 | raise Exception('Task dependencies only work with Workflows') 150 | if task.name in self._dependencies: 151 | raise Exception('`{}` already depends on `{}`'.format(self.name, task.name)) 152 | if self.name == task.name: 153 | raise Exception('A task cannot depend on itself') 154 | self._dependencies.add(task.name) 155 | 156 | def get_new_instance(self, 157 | scheduled=False, 158 | status='queued', 159 | workflow_instance_id=None, 160 | run_at=None, 161 | priority=None, 162 | max_attempts=None, 163 | timeout=None, 164 | retry_delay=None, 165 | unique=None): 166 | return TaskInstance( 167 | task_name=self.name, 168 | workflow_instance_id=workflow_instance_id, 169 | scheduled=scheduled, 170 | push=self.push_destination != None, 171 | status=status, 172 | priority=priority or self.default_priority, 173 | run_at=run_at or datetime.utcnow(), 174 | max_attempts=max_attempts or (self.retries + 1), 175 | timeout=timeout or self.timeout, 176 | retry_delay=retry_delay or self.retry_delay, 177 | unique=unique) 178 | 179 | def execute(self, task_instance): 180 | raise NotImplementedError() 181 | 182 | def on_kill(self): 183 | pass 184 | 185 | pull_sql = """ 186 | WITH nextTasks as ( 187 | SELECT id, status, started_at 188 | FROM task_instances 189 | WHERE 190 | {} 191 | run_at <= :now AND 192 | attempts < max_attempts AND 193 | (status = 'queued' OR 194 | (status = 'running' AND (:now > (locked_at + INTERVAL '1 second' * timeout))) OR 195 | (status = 'retry' AND (:now > (locked_at + INTERVAL '1 second' * retry_delay)))) 196 | ORDER BY 197 | CASE WHEN priority = 'critical' 198 | THEN 1 199 | WHEN priority = 'high' 200 | THEN 2 201 | WHEN priority = 'normal' 202 | THEN 3 203 | WHEN priority = 'low' 204 | THEN 4 205 | END, 206 | run_at 207 | LIMIT :max_tasks 208 | FOR UPDATE SKIP LOCKED 209 | ) 210 | UPDATE task_instances SET 211 | status = 'running'::taskflow_statuses, 212 | worker_id = :worker_id, 213 | locked_at = :now, 214 | started_at = COALESCE(nextTasks.started_at, :now), 215 | attempts = attempts + 1 216 | FROM nextTasks 217 | WHERE task_instances.id = nextTasks.id 218 | RETURNING task_instances.*; 219 | """ 220 | 221 | task_names_filter = '\n task_instances.name = ANY(:task_names)\n AND' 222 | push_filter = '\n task_instances.push = true\n AND' 223 | 224 | class Taskflow(object): 225 | def __init__(self, monitoring=None): 226 | self._workflows = dict() 227 | self._tasks = dict() 228 | self._push_workers = dict() 229 | 230 | self.monitoring = monitoring or Monitor() 231 | 232 | def set_monitoring(self, monitoring): 233 | self.monitoring = monitoring 234 | 235 | def add_workflow(self, workflow): 236 | self._workflows[workflow.name] = workflow 237 | 238 | def add_workflows(self, workflows): 239 | for workflow in workflows: 240 | self.add_workflow(workflow) 241 | 242 | def get_workflow(self, workflow_name): 243 | return self._workflows[workflow_name] 244 | 245 | def get_workflows(self): 246 | return self._workflows.values() 247 | 248 | def add_task(self, task): 249 | if task.workflow != None: 250 | raise Exception('Tasks with workflows are not added individually, just add the workflow') 251 | self._tasks[task.name] = task 252 | 253 | def add_tasks(self, tasks): 254 | for task in tasks: 255 | self.add_task(task) 256 | 257 | def get_task(self, task_name): 258 | if task_name in self._tasks: 259 | return self._tasks[task_name] 260 | for workflow in self._workflows.values(): 261 | task = workflow.get_task(task_name) 262 | if task: 263 | return task 264 | 265 | def get_tasks(self): 266 | return self._tasks.values() 267 | 268 | def add_push_worker(self, push_worker): 269 | self._push_workers[push_worker.push_type] = push_worker 270 | 271 | def get_push_worker(self, push_type): 272 | if push_type not in self._push_workers: 273 | return None 274 | return self._push_workers[push_type] 275 | 276 | def sync_tasks(self, session, tasks): 277 | for task in tasks: 278 | existing = session.query(Task).filter(Task.name == task.name).one_or_none() 279 | if existing: 280 | task.active = existing.active 281 | session.merge(task) 282 | else: 283 | session.add(task) 284 | 285 | def sync_db(self, session, read_only=False): 286 | for workflow_name in self._workflows: 287 | workflow = self._workflows[workflow_name] 288 | existing = session.query(Workflow).filter(Workflow.name == workflow_name).one_or_none() 289 | if existing: 290 | workflow.active = existing.active 291 | session.merge(workflow) 292 | else: 293 | session.add(workflow) 294 | 295 | self.sync_tasks(session, workflow.get_tasks()) 296 | 297 | self.sync_tasks(session, self._tasks.values()) 298 | 299 | if not read_only: 300 | session.commit() 301 | 302 | def pull(self, session, worker_id, task_names=None, max_tasks=1, now=None, push=False): 303 | if now == None: 304 | now = datetime.utcnow() 305 | 306 | params = { 307 | 'worker_id': worker_id, 308 | 'now': now, 309 | 'max_tasks': max_tasks 310 | } 311 | 312 | filters = '' 313 | if task_names != None: 314 | filters += task_names_filter 315 | params['task_names'] = task_names 316 | if push: 317 | filters += push_filter 318 | 319 | pull_sql_with_filters = pull_sql.format(filters) 320 | 321 | task_instances = session.query(TaskInstance)\ 322 | .from_statement(text(pull_sql_with_filters))\ 323 | .params(**params)\ 324 | .all() 325 | 326 | return task_instances 327 | 328 | class SchedulableInstance(BaseModel): 329 | __abstract__ = True 330 | 331 | id = Column(BigInteger, primary_key=True) 332 | scheduled = Column(Boolean, nullable=False, default=False) 333 | run_at = Column(DateTime, nullable=False, default=datetime.utcnow) 334 | started_at = Column(DateTime) 335 | ended_at = Column(DateTime) 336 | status = Column(Enum('queued', 337 | 'pushed', 338 | 'running', 339 | 'retry', 340 | 'dequeued', 341 | 'failed', 342 | 'success', 343 | name='taskflow_statuses'), 344 | nullable=False, 345 | default='queued') 346 | priority = Column(Enum('critical', 347 | 'high', 348 | 'normal', 349 | 'low', 350 | name='taskflow_priorities'), 351 | nullable=False, 352 | default='normal') 353 | unique = Column(String) 354 | created_at = Column(DateTime, 355 | nullable=False, 356 | server_default=func.now()) 357 | updated_at = Column(DateTime, 358 | nullable=False, 359 | server_default=func.now(), 360 | onupdate=func.now()) 361 | 362 | def complete(self, session, taskflow, status, now=None): 363 | if now == None: 364 | now = datetime.utcnow() 365 | 366 | self.status = status 367 | self.ended_at = now 368 | 369 | session.commit() 370 | 371 | if status == 'success': 372 | taskflow.monitoring.task_success(session, self) 373 | else: 374 | taskflow.monitoring.task_failed(session, self) 375 | 376 | def succeed(self, session, taskflow, now=None): 377 | self.complete(session, taskflow, 'success', now=now) 378 | 379 | def fail(self, session, taskflow, now=None, dry_run=False): 380 | if now == None: 381 | now = datetime.utcnow() 382 | 383 | if isinstance(self, TaskInstance) and self.attempts < self.max_attempts: 384 | self.status = 'retry' 385 | self.locked_at = now 386 | 387 | if not dry_run: 388 | session.commit() 389 | taskflow.monitoring.task_retry(session, self) 390 | else: 391 | self.complete(session, taskflow, 'failed', now=now) 392 | 393 | @listens_for(SchedulableInstance, 'instrument_class', propagate=True) 394 | def receive_mapper_configured(mapper, class_): 395 | class_.build_indexes() 396 | 397 | class WorkflowInstance(SchedulableInstance): 398 | __tablename__ = 'workflow_instances' 399 | 400 | workflow_name = Column(String, nullable=False) 401 | params = Column(JSONB) 402 | task_instances = relationship('TaskInstance', backref='workflow', passive_deletes=True) 403 | 404 | def __repr__(self): 405 | return ''.format( 406 | self.id, 407 | self.workflow_name, 408 | self.run_at, 409 | self.status) 410 | 411 | @classmethod 412 | def build_indexes(cls): 413 | Index('index_unique_workflow', 414 | cls.workflow_name, 415 | cls.unique, 416 | unique=True, 417 | postgresql_where= 418 | cls.status.in_(['queued','pushed','running','retry'])) 419 | 420 | class TaskInstance(SchedulableInstance): 421 | __tablename__ = 'task_instances' 422 | 423 | task_name = Column(String, nullable=False) 424 | workflow_instance_id = Column(BigInteger, ForeignKey('workflow_instances.id', ondelete='CASCADE')) 425 | push = Column(Boolean, nullable=False) 426 | locked_at = Column(DateTime) ## TODO: should workflow instaces have locked_at as well ? 427 | worker_id = Column(String) 428 | params = Column(JSONB, default={}) 429 | push_state = Column(JSONB) 430 | attempts = Column(Integer, nullable=False, default=0) 431 | max_attempts = Column(Integer, nullable=False, default=1) 432 | timeout = Column(Integer, nullable=False) 433 | retry_delay = Column(Integer, nullable=False) 434 | 435 | def __repr__(self): 436 | return ''.format( 437 | self.id, 438 | self.task_name, 439 | self.workflow_instance_id, 440 | self.status) 441 | 442 | @classmethod 443 | def build_indexes(cls): 444 | Index('index_unique_task', 445 | cls.task_name, 446 | cls.unique, 447 | unique=True, 448 | postgresql_where= 449 | cls.status.in_(['queued','pushed','running','retry'])) 450 | 451 | class TaskflowEvent(BaseModel): 452 | __tablename__ = 'taskflow_events' 453 | 454 | id = Column(BigInteger, primary_key=True) 455 | workflow_instance = Column(BigInteger, ForeignKey('workflow_instances.id')) 456 | task_instance = Column(BigInteger, ForeignKey('task_instances.id')) 457 | timestamp = Column(DateTime, nullable=False) 458 | event = Column(String, nullable=False) 459 | message = Column(String) 460 | 461 | class User(UserAuthMixin, UserMixin, BaseModel): 462 | __tablename__ = 'users' 463 | 464 | id = Column(Integer, primary_key=True) 465 | active = Column(Boolean, nullable=False) 466 | email = Column(String) 467 | ## username from restful-ben 468 | ## hashed_password from restful-ben 469 | ## password property from restful-ben 470 | role = Column(Enum('normal','admin', name='user_roles'), nullable=False) 471 | created_at = Column(DateTime, 472 | nullable=False, 473 | server_default=func.now()) 474 | updated_at = Column(DateTime, 475 | nullable=False, 476 | server_default=func.now(), 477 | onupdate=func.now()) 478 | 479 | @property 480 | def is_active(self): 481 | return self.active 482 | 483 | def __repr__(self): 484 | return ''.format(self.id, \ 485 | self.active, \ 486 | self.username, \ 487 | self.email) 488 | 489 | -------------------------------------------------------------------------------- /taskflow/core/pusher.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | from itertools import groupby 3 | import logging 4 | 5 | from .models import Workflow, WorkflowInstance, TaskInstance 6 | 7 | class Pusher(object): 8 | def __init__(self, taskflow, dry_run=None, now_override=None): 9 | self.logger = logging.getLogger('Pusher') 10 | 11 | self.taskflow = taskflow 12 | 13 | self.dry_run = dry_run 14 | self.now_override = now_override 15 | 16 | def now(self): 17 | """Allows for dry runs and tests to use a specific datetime as now""" 18 | if self.now_override: 19 | return self.now_override 20 | return datetime.utcnow() 21 | 22 | def get_push_destination(self, task_instance): 23 | return self.taskflow.get_task(task_instance.task_name).push_destination 24 | 25 | def sync_task_states(self, session): 26 | ## TODO: paginate? 27 | task_instances = session.query(TaskInstance)\ 28 | .filter(TaskInstance.push == True, TaskInstance.status.in_(['pushed','running'])).all() 29 | 30 | for push_destination, task_instances in groupby(task_instances, self.get_push_destination): 31 | self.logger.info('Syncing states with %s', push_destination) 32 | 33 | try: 34 | push_worker = self.taskflow.get_push_worker(push_destination) 35 | push_worker.sync_task_instance_states(session, self.dry_run, task_instances, self.now()) 36 | except Exception: 37 | ## TODO: rollback? 38 | self.logger.exception('Exception syncing with %s', push_destination) 39 | 40 | 41 | def push_queued_task_instances(self, session): 42 | task_instances = self.taskflow.pull(session, 'Pusher', max_tasks=100, now=self.now(), push=True) 43 | 44 | for push_destination, task_instances in groupby(task_instances, self.get_push_destination): 45 | self.logger.info('Pushing to %s', push_destination) 46 | 47 | try: 48 | push_worker = self.taskflow.get_push_worker(push_destination) 49 | push_worker.push_task_instances(session, self.dry_run, task_instances) 50 | except Exception: 51 | ## TODO: rollback? 52 | self.logger.exception('Exception pushing to %s', push_destination) 53 | 54 | def run(self, session): 55 | self.logger.info('*** Starting Pusher Run ***') 56 | 57 | self.logger.info('Pushing queued task instances') 58 | self.push_queued_task_instances(session) 59 | 60 | self.logger.info('Syncing pushed task instance states') 61 | self.sync_task_states(session) 62 | 63 | self.logger.info('*** End Pusher Run ***') 64 | -------------------------------------------------------------------------------- /taskflow/core/scheduler.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | import logging 3 | 4 | from toposort import toposort 5 | from sqlalchemy import or_, and_ 6 | 7 | from .models import Workflow, WorkflowInstance, Task, TaskInstance 8 | 9 | class Scheduler(object): 10 | def __init__(self, taskflow, dry_run=False, now_override=None): 11 | self.logger = logging.getLogger('Scheduler') 12 | 13 | self.taskflow = taskflow 14 | 15 | self.dry_run = dry_run 16 | self.now_override = now_override 17 | 18 | def now(self): 19 | """Allows for dry runs and tests to use a specific datetime as now""" 20 | if self.now_override: 21 | return self.now_override 22 | return datetime.utcnow() 23 | 24 | def queue_task(self, session, task, run_at): 25 | if run_at == None: 26 | run_at = self.now() 27 | 28 | task_instance = task.get_new_instance(scheduled=True, 29 | run_at=run_at, 30 | unique='scheduled_' + run_at.isoformat()) 31 | 32 | self.logger.info('Queuing task: %s %s', task.name, run_at) 33 | 34 | if not self.dry_run: 35 | session.add(task_instance) 36 | session.commit() 37 | 38 | def queue_workflow_task(self, session, workflow, task_name, workflow_instance, run_at=None): 39 | if run_at == None: 40 | run_at = self.now() 41 | 42 | task = workflow.get_task(task_name) 43 | 44 | task_instance = task.get_new_instance( 45 | scheduled=True, 46 | run_at=run_at, 47 | workflow_instance_id=workflow_instance.id, 48 | priority=workflow_instance.priority or workflow.default_priority, 49 | unique='scheduled_' + run_at.isoformat()) 50 | 51 | self.logger.info('Queuing workflow task: %s %s %s', workflow.name, task.name, run_at) 52 | 53 | if not self.dry_run: 54 | session.add(task_instance) 55 | 56 | def queue_workflow_tasks(self, session, workflow_instance): 57 | workflow = self.taskflow.get_workflow(workflow_instance.workflow_name) 58 | dep_graph = workflow.get_dependencies_graph() 59 | dep_graph = list(toposort(dep_graph)) 60 | 61 | results = session.query(TaskInstance)\ 62 | .filter(TaskInstance.workflow_instance_id == workflow_instance.id).all() 63 | workflow_task_instances = dict() 64 | for instance in results: 65 | workflow_task_instances[instance.task_name] = instance 66 | 67 | ## dep_graph looks like [{'task2', 'task1'}, {'task3'}, {'task4'}] 68 | ## a list of sets where each set is a parallel step 69 | total_complete_steps = 0 70 | failed = False 71 | for step in dep_graph: 72 | total_in_step = len(step) 73 | total_complete = 0 74 | tasks_to_queue = [] 75 | 76 | for task_name in step: 77 | if task_name in workflow_task_instances: 78 | if workflow_task_instances[task_name].status == 'success': 79 | total_complete += 1 80 | elif workflow_task_instances[task_name].status == 'failed': 81 | failed = True 82 | break 83 | else: 84 | tasks_to_queue.append(task_name) 85 | 86 | if failed: 87 | break 88 | 89 | if not self.dry_run: 90 | for task_name in tasks_to_queue: 91 | self.queue_workflow_task(session, workflow, task_name, workflow_instance) 92 | 93 | if len(tasks_to_queue) > 0 and total_complete == total_in_step: 94 | raise Exception('Attempting to queue tasks for a completed workflow step') 95 | 96 | if total_complete < total_in_step: 97 | break 98 | else: 99 | total_complete_steps += 1 100 | 101 | if failed: 102 | workflow_instance.status = 'failed' 103 | workflow_instance.ended_at = self.now() 104 | self.logger.info('Workflow {} - {} failed'.format(workflow_instance.workflow_name, workflow_instance.id)) 105 | if not self.dry_run: 106 | session.commit() 107 | self.taskflow.monitoring.workflow_failed(session, workflow_instance) 108 | elif total_complete_steps == len(dep_graph): 109 | workflow_instance.status = 'success' 110 | workflow_instance.ended_at = self.now() 111 | self.logger.info('Workflow {} - {} succeeded'.format(workflow_instance.workflow_name, workflow_instance.id)) 112 | if not self.dry_run: 113 | session.commit() 114 | self.taskflow.monitoring.workflow_success(session, workflow_instance) 115 | 116 | def queue_workflow(self, session, workflow, run_at): 117 | workflow_instance = workflow.get_new_instance( 118 | scheduled=True, 119 | run_at=run_at, 120 | unique='scheduled_' + run_at.isoformat()) 121 | 122 | self.logger.info('Queuing workflow: %s', workflow.name) 123 | 124 | if not self.dry_run: 125 | session.add(workflow_instance) 126 | 127 | if workflow_instance.run_at <= self.now(): 128 | self.queue_workflow_tasks(session, workflow_instance) 129 | 130 | if not self.dry_run: 131 | session.commit() 132 | 133 | def schedule_recurring(self, session, definition_class): 134 | """Schedules recurring Workflows or Tasks 135 | definition_class - Workflow or Task""" 136 | 137 | if definition_class == Workflow: 138 | instance_class = WorkflowInstance 139 | recurring_items = self.taskflow.get_workflows() 140 | elif definition_class == Task: 141 | instance_class = TaskInstance 142 | recurring_items = self.taskflow.get_tasks() 143 | else: 144 | raise Exception('definition_class must be Workflow or Task') 145 | 146 | now = self.now() 147 | 148 | ## get Workflows or Tasks from Taskflow instance 149 | recurring_items = filter(lambda item: item.active == True and item.schedule != None, 150 | recurring_items) 151 | 152 | for item in recurring_items: 153 | self.logger.info('Scheduling recurring %s: %s', definition_class.__name__.lower(), item.name) 154 | 155 | try: 156 | if definition_class == Workflow: 157 | filters = (instance_class.workflow_name == item.name,) 158 | else: 159 | filters = (instance_class.task_name == item.name,) 160 | filters += (instance_class.scheduled == True,) 161 | 162 | ## Get the most recent instance of the recurring item 163 | ## TODO: order by started_at instead ? 164 | most_recent_instance = session.query(instance_class)\ 165 | .filter(*filters)\ 166 | .order_by(instance_class.run_at.desc())\ 167 | .first() 168 | 169 | if not most_recent_instance or most_recent_instance.status in ['success','failed']: 170 | if not most_recent_instance: ## first run 171 | next_run = item.next_run(base_time=now) 172 | else: 173 | next_run = item.next_run(base_time=most_recent_instance.run_at) 174 | last_run = item.last_run(base_time=now) 175 | if last_run > next_run: 176 | next_run = last_run 177 | 178 | if item.start_date and next_run < item.start_date or \ 179 | item.end_date and next_run > item.end_date: 180 | self.logger.info('%s is not within its scheduled range', item.name) 181 | continue 182 | 183 | if definition_class == Workflow: 184 | self.queue_workflow(session, item, next_run) 185 | else: 186 | self.queue_task(session, item, next_run) 187 | except Exception: 188 | self.logger.exception('Exception scheduling %s', item.name) 189 | session.rollback() 190 | 191 | def advance_workflows_forward(self, session): 192 | """Moves queued and running workflows forward""" 193 | now = self.now() 194 | 195 | queued_running_workflow_instances = \ 196 | session.query(WorkflowInstance)\ 197 | .filter(or_(WorkflowInstance.status == 'running', 198 | and_(WorkflowInstance.status == 'queued', 199 | WorkflowInstance.run_at <= now)))\ 200 | .all() ## TODO: paginate? 201 | 202 | for workflow_instance in queued_running_workflow_instances: 203 | self.logger.info('Checking %s - %s for advancement', workflow_instance.workflow_name, workflow_instance.id) 204 | try: 205 | if workflow_instance.status == 'queued': 206 | workflow_instance.status = 'running' 207 | workflow_instance.started_at = self.now() 208 | self.logger.info('Starting workflow {} - {} failed'.format(workflow_instance.workflow_name, workflow_instance.id)) 209 | self.queue_workflow_tasks(session, workflow_instance) 210 | if not self.dry_run: 211 | session.commit() 212 | elif workflow_instance.status == 'running': 213 | ## TODO: timeout queued workflow instances that have gone an interval past their run_at 214 | self.queue_workflow_tasks(session, workflow_instance) 215 | if not self.dry_run: 216 | session.commit() 217 | except Exception: 218 | self.logger.exception('Exception scheduling %s', workflow_instance.workflow_name) 219 | session.rollback() 220 | 221 | def fail_timedout_task_instances(self, session): 222 | ## TODO: return info using RETURNING and log 223 | if not self.dry_run: 224 | session.execute( 225 | "UPDATE task_instances SET status = 'failed', ended_at = :now " + 226 | "WHERE status in ('running','retry') AND " + 227 | "(:now > (locked_at + INTERVAL '1 second' * timeout)) AND " + 228 | "attempts >= max_attempts", {'now': self.now()}) 229 | 230 | def run(self, session): 231 | ## TODO: what happens when a schedule changes? - Looks like it has to what until the current future runs 232 | 233 | self.logger.info('*** Starting Scheduler Run ***') 234 | 235 | ##### Workflow scheduling 236 | 237 | self.logger.info('Scheduling recurring workflows') 238 | 239 | self.schedule_recurring(session, Workflow) 240 | 241 | self.logger.info('Advancing workflows') 242 | 243 | self.advance_workflows_forward(session) 244 | 245 | ## TODO: start / advance non recurring workflows 246 | 247 | 248 | ##### Task scheduling - tasks that do not belong to a workflow 249 | 250 | self.logger.info('Scheduling recurring tasks') 251 | 252 | self.schedule_recurring(session, Task) 253 | 254 | self.logger.info('Failing timed out tasks') 255 | 256 | self.fail_timedout_task_instances(session) 257 | 258 | if not self.dry_run: 259 | self.logger.info('Sending Heartbeat') 260 | self.taskflow.monitoring.heartbeat_scheduler(session) 261 | 262 | self.logger.info('*** End Scheduler Run ***') 263 | -------------------------------------------------------------------------------- /taskflow/core/worker.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import signal 3 | import sys 4 | 5 | from .models import Workflow, WorkflowInstance, Task, TaskInstance 6 | 7 | class Worker(object): 8 | def __init__(self, taskflow): 9 | self.logger = logging.getLogger('Worker') 10 | self.taskflow = taskflow 11 | 12 | signal.signal(signal.SIGINT, self.on_kill) 13 | signal.signal(signal.SIGTERM, self.on_kill) 14 | 15 | def on_kill(self, sig_num, stack_frame): 16 | if hasattr(self, 'task'): 17 | self.task.on_kill() 18 | sys.exit(0) 19 | 20 | def execute(self, session, task_instance): 21 | try: 22 | task = self.taskflow.get_task(task_instance.task_name) 23 | if not task: 24 | raise Exception('Task `{}` does not exist'.format(task_instance.task_name)) 25 | 26 | self.current_task = task 27 | task.execute(task_instance) 28 | except Exception: 29 | self.logger.exception('Error executing: %s %s', task_instance.task_name, task_instance.id) 30 | task_instance.fail(session, self.taskflow) 31 | return False 32 | task_instance.succeed(session, self.taskflow) 33 | return True 34 | -------------------------------------------------------------------------------- /taskflow/db.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from alembic import command 4 | from alembic.config import Config 5 | from alembic.migration import MigrationContext 6 | from sqlalchemy import create_engine 7 | 8 | from .core.models import BaseModel 9 | 10 | def get_alembic_config(connection_string): 11 | current_dir = os.path.dirname(os.path.abspath(__file__)) 12 | scripts_directory = os.path.join(current_dir, 'migrations') 13 | alembic_config = Config(os.path.join(current_dir, 'alembic.ini')) 14 | alembic_config.set_main_option('script_location', scripts_directory) 15 | alembic_config.set_main_option('sqlalchemy.url', connection_string) 16 | return alembic_config 17 | 18 | def migrate_db(connection_string): 19 | alembic_config = get_alembic_config(connection_string) 20 | command.upgrade(alembic_config, 'heads') 21 | 22 | def init_db(connection_string): 23 | engine = create_engine(connection_string) 24 | 25 | BaseModel.metadata.create_all(engine) 26 | 27 | alembic_config = get_alembic_config(connection_string) 28 | command.stamp(alembic_config, 'head') 29 | -------------------------------------------------------------------------------- /taskflow/migrations/README: -------------------------------------------------------------------------------- 1 | Generic single-database configuration. -------------------------------------------------------------------------------- /taskflow/migrations/env.py: -------------------------------------------------------------------------------- 1 | from __future__ import with_statement 2 | from alembic import context 3 | from sqlalchemy import engine_from_config, pool 4 | from logging.config import fileConfig 5 | 6 | # this is the Alembic Config object, which provides 7 | # access to the values within the .ini file in use. 8 | config = context.config 9 | 10 | # Interpret the config file for Python logging. 11 | # This line sets up loggers basically. 12 | fileConfig(config.config_file_name) 13 | 14 | # add your model's MetaData object here 15 | # for 'autogenerate' support 16 | # from myapp import mymodel 17 | # target_metadata = mymodel.Base.metadata 18 | target_metadata = None 19 | 20 | # other values from the config, defined by the needs of env.py, 21 | # can be acquired: 22 | # my_important_option = config.get_main_option("my_important_option") 23 | # ... etc. 24 | 25 | 26 | def run_migrations_offline(): 27 | """Run migrations in 'offline' mode. 28 | 29 | This configures the context with just a URL 30 | and not an Engine, though an Engine is acceptable 31 | here as well. By skipping the Engine creation 32 | we don't even need a DBAPI to be available. 33 | 34 | Calls to context.execute() here emit the given string to the 35 | script output. 36 | 37 | """ 38 | url = config.get_main_option("sqlalchemy.url") 39 | context.configure( 40 | url=url, target_metadata=target_metadata, literal_binds=True) 41 | 42 | with context.begin_transaction(): 43 | context.run_migrations() 44 | 45 | 46 | def run_migrations_online(): 47 | """Run migrations in 'online' mode. 48 | 49 | In this scenario we need to create an Engine 50 | and associate a connection with the context. 51 | 52 | """ 53 | connectable = engine_from_config( 54 | config.get_section(config.config_ini_section), 55 | prefix='sqlalchemy.', 56 | poolclass=pool.NullPool) 57 | 58 | with connectable.connect() as connection: 59 | context.configure( 60 | connection=connection, 61 | target_metadata=target_metadata 62 | ) 63 | 64 | with context.begin_transaction(): 65 | context.run_migrations() 66 | 67 | if context.is_offline_mode(): 68 | run_migrations_offline() 69 | else: 70 | run_migrations_online() 71 | -------------------------------------------------------------------------------- /taskflow/migrations/script.py.mako: -------------------------------------------------------------------------------- 1 | """${message} 2 | 3 | Revision ID: ${up_revision} 4 | Revises: ${down_revision | comma,n} 5 | Create Date: ${create_date} 6 | 7 | """ 8 | from alembic import op 9 | import sqlalchemy as sa 10 | ${imports if imports else ""} 11 | 12 | # revision identifiers, used by Alembic. 13 | revision = ${repr(up_revision)} 14 | down_revision = ${repr(down_revision)} 15 | branch_labels = ${repr(branch_labels)} 16 | depends_on = ${repr(depends_on)} 17 | 18 | 19 | def upgrade(): 20 | ${upgrades if upgrades else "pass"} 21 | 22 | 23 | def downgrade(): 24 | ${downgrades if downgrades else "pass"} 25 | -------------------------------------------------------------------------------- /taskflow/migrations/versions/1e64f390802b_add_unique_task_and_workflow_indexes.py: -------------------------------------------------------------------------------- 1 | """add unique task and workflow indexes 2 | 3 | Revision ID: 1e64f390802b 4 | Revises: 6ef79b56ad4a 5 | Create Date: 2017-06-19 17:02:25.402708+00:00 6 | 7 | """ 8 | from alembic import op 9 | import sqlalchemy as sa 10 | 11 | 12 | # revision identifiers, used by Alembic. 13 | revision = '1e64f390802b' 14 | down_revision = '6ef79b56ad4a' 15 | branch_labels = None 16 | depends_on = None 17 | 18 | 19 | def upgrade(): 20 | op.execute("CREATE UNIQUE INDEX index_unique_workflow ON workflow_instances USING btree (workflow_name, \"unique\") WHERE status = ANY (ARRAY['queued'::taskflow_statuses, 'pushed'::taskflow_statuses, 'running'::taskflow_statuses, 'retry'::taskflow_statuses]);") 21 | op.execute("CREATE UNIQUE INDEX index_unique_task ON task_instances USING btree (task_name, \"unique\") WHERE status = ANY (ARRAY['queued'::taskflow_statuses, 'pushed'::taskflow_statuses, 'running'::taskflow_statuses, 'retry'::taskflow_statuses]);") 22 | 23 | def downgrade(): 24 | op.execute('DROP INDEX index_unique_workflow;') 25 | op.execute('DROP INDEX index_unique_task;') 26 | -------------------------------------------------------------------------------- /taskflow/migrations/versions/6ef79b56ad4a_rename_status_retry.py: -------------------------------------------------------------------------------- 1 | """rename status retry 2 | 3 | Revision ID: 6ef79b56ad4a 4 | Revises: 94e7b91d83d5 5 | Create Date: 2017-06-19 15:06:50.441524+00:00 6 | 7 | """ 8 | from alembic import op 9 | import sqlalchemy as sa 10 | 11 | 12 | # revision identifiers, used by Alembic. 13 | revision = '6ef79b56ad4a' 14 | down_revision = '94e7b91d83d5' 15 | branch_labels = None 16 | depends_on = None 17 | 18 | ## ref http://blog.yo1.dog/updating-enum-values-in-postgresql-the-safe-and-easy-way/ 19 | 20 | def upgrade(): 21 | op.execute(""" 22 | BEGIN; 23 | ALTER TYPE taskflow_statuses RENAME TO taskflow_statuses_old; 24 | CREATE TYPE taskflow_statuses AS ENUM('queued','pushed','running','retry','dequeued','failed','success'); 25 | CREATE TYPE taskflow_statuses_inter AS ENUM('queued','pushed','running','retry','retrying','dequeued','failed','success'); 26 | 27 | ALTER TABLE workflow_instances ALTER COLUMN status TYPE taskflow_statuses_inter USING status::text::taskflow_statuses_inter; 28 | UPDATE workflow_instances SET status = 'retry' WHERE status = 'retrying'; 29 | ALTER TABLE workflow_instances ALTER COLUMN status TYPE taskflow_statuses USING status::text::taskflow_statuses; 30 | ALTER TABLE task_instances ALTER COLUMN status TYPE taskflow_statuses_inter USING status::text::taskflow_statuses_inter; 31 | UPDATE task_instances SET status = 'retry' WHERE status = 'retrying'; 32 | ALTER TABLE task_instances ALTER COLUMN status TYPE taskflow_statuses USING status::text::taskflow_statuses; 33 | 34 | DROP TYPE taskflow_statuses_old; 35 | DROP TYPE taskflow_statuses_inter; 36 | COMMIT; 37 | """) 38 | 39 | def downgrade(): 40 | op.execute(""" 41 | BEGIN; 42 | ALTER TYPE taskflow_statuses RENAME TO taskflow_statuses_old; 43 | CREATE TYPE taskflow_statuses AS ENUM('queued','pushed','running','retrying','dequeued','failed','success'); 44 | CREATE TYPE taskflow_statuses_inter AS ENUM('queued','pushed','running','retry','retrying','dequeued','failed','success'); 45 | 46 | ALTER TABLE workflow_instances ALTER COLUMN status TYPE taskflow_statuses_inter USING status::text::taskflow_statuses_inter; 47 | UPDATE workflow_instances SET status = 'retrying' WHERE status = 'retry'; 48 | ALTER TABLE workflow_instances ALTER COLUMN status TYPE taskflow_statuses USING status::text::taskflow_statuses; 49 | ALTER TABLE task_instances ALTER COLUMN status TYPE taskflow_statuses_inter USING status::text::taskflow_statuses_inter; 50 | UPDATE task_instances SET status = 'retrying' WHERE status = 'retry'; 51 | ALTER TABLE task_instances ALTER COLUMN status TYPE taskflow_statuses USING status::text::taskflow_statuses; 52 | 53 | DROP TYPE taskflow_statuses_old; 54 | DROP TYPE taskflow_statuses_inter; 55 | COMMIT; 56 | """) 57 | -------------------------------------------------------------------------------- /taskflow/migrations/versions/94e7b91d83d5_create_task_workflow_name_column.py: -------------------------------------------------------------------------------- 1 | """Create Task workflow_name column 2 | 3 | Revision ID: 94e7b91d83d5 4 | Revises: 5 | Create Date: 2017-06-14 02:33:42.738473+00:00 6 | 7 | """ 8 | from alembic import op 9 | import sqlalchemy as sa 10 | 11 | 12 | # revision identifiers, used by Alembic. 13 | revision = '94e7b91d83d5' 14 | down_revision = None 15 | branch_labels = None 16 | depends_on = None 17 | 18 | 19 | def upgrade(): 20 | op.add_column('tasks', 21 | sa.Column('workflow_name', sa.String, sa.ForeignKey('workflows.name')) 22 | ) 23 | 24 | 25 | def downgrade(): 26 | op.drop_column('tasks', 'workflow_name') 27 | -------------------------------------------------------------------------------- /taskflow/monitoring/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CityOfPhiladelphia/taskflow/5c86b50c6d8bce43a04fba1b856b1d56f3c36b6a/taskflow/monitoring/__init__.py -------------------------------------------------------------------------------- /taskflow/monitoring/aws.py: -------------------------------------------------------------------------------- 1 | import boto3 2 | 3 | from .base import MonitorDestination 4 | 5 | class AWSMonitor(MonitorDestination): 6 | def __init__(self, 7 | metric_prefix='', 8 | metric_namespace='taskflow', 9 | *args, **kwargs): 10 | self.metric_namespace = metric_namespace 11 | self.metric_prefix = metric_prefix 12 | self.cloudwatch = boto3.client('cloudwatch') 13 | 14 | super(AWSMonitor, self).__init__(*args, **kwargs) 15 | 16 | def heartbeat_scheduler(self, session): 17 | self.cloudwatch.put_metric_data( 18 | Namespace=self.metric_namespace, 19 | MetricData=[ 20 | { 21 | 'MetricName': self.metric_prefix + 'scheduler_heartbeat', 22 | 'Value': 1, 23 | 'Unit': 'Count' 24 | } 25 | ]) 26 | 27 | def task_retry(self, session, task_instance): 28 | self.cloudwatch.put_metric_data( 29 | Namespace=self.metric_namespace, 30 | MetricData=[ 31 | { 32 | 'MetricName': self.metric_prefix + 'task_retry', 33 | 'Value': 1, 34 | 'Unit': 'Count' 35 | }, 36 | { 37 | 'MetricName': self.metric_prefix + 'task_retry', 38 | 'Dimensions': [ 39 | { 40 | 'Name': 'task_name', 41 | 'Value': task_instance.task_name 42 | } 43 | ], 44 | 'Value': 1, 45 | 'Unit': 'Count' 46 | } 47 | ]) 48 | 49 | def task_failed(self, session, task_instance): 50 | self.cloudwatch.put_metric_data( 51 | Namespace=self.metric_namespace, 52 | MetricData=[ 53 | { 54 | 'MetricName': self.metric_prefix + 'task_failure', 55 | 'Value': 1, 56 | 'Unit': 'Count' 57 | }, 58 | { 59 | 'MetricName': self.metric_prefix + 'task_failure', 60 | 'Dimensions': [ 61 | { 62 | 'Name': 'task_name', 63 | 'Value': task_instance.task_name 64 | } 65 | ], 66 | 'Value': 1, 67 | 'Unit': 'Count' 68 | } 69 | ]) 70 | 71 | def task_success(self, session, task_instance): 72 | self.cloudwatch.put_metric_data( 73 | Namespace=self.metric_namespace, 74 | MetricData=[ 75 | { 76 | 'MetricName': self.metric_prefix + 'task_success', 77 | 'Value': 1, 78 | 'Unit': 'Count' 79 | }, 80 | { 81 | 'MetricName': self.metric_prefix + 'task_success', 82 | 'Dimensions': [ 83 | { 84 | 'Name': 'task_name', 85 | 'Value': task_instance.task_name 86 | } 87 | ], 88 | 'Value': 1, 89 | 'Unit': 'Count' 90 | } 91 | ]) 92 | 93 | def workflow_failed(self, session, workflow_instance): 94 | self.cloudwatch.put_metric_data( 95 | Namespace=self.metric_namespace, 96 | MetricData=[ 97 | { 98 | 'MetricName': self.metric_prefix + 'workflow_failure', 99 | 'Value': 1, 100 | 'Unit': 'Count' 101 | }, 102 | { 103 | 'MetricName': self.metric_prefix + 'workflow_failure', 104 | 'Dimensions': [ 105 | { 106 | 'Name': 'workflow_name', 107 | 'Value': workflow_instance.workflow_name 108 | } 109 | ], 110 | 'Value': 1, 111 | 'Unit': 'Count' 112 | } 113 | ]) 114 | 115 | def workflow_success(self, session, workflow_instance): 116 | self.cloudwatch.put_metric_data( 117 | Namespace=self.metric_namespace, 118 | MetricData=[ 119 | { 120 | 'MetricName': self.metric_prefix + 'workflow_success', 121 | 'Value': 1, 122 | 'Unit': 'Count' 123 | }, 124 | { 125 | 'MetricName': self.metric_prefix + 'workflow_success', 126 | 'Dimensions': [ 127 | { 128 | 'Name': 'workflow_name', 129 | 'Value': workflow_instance.workflow_name 130 | } 131 | ], 132 | 'Value': 1, 133 | 'Unit': 'Count' 134 | } 135 | ]) 136 | -------------------------------------------------------------------------------- /taskflow/monitoring/base.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | class Monitor(object): 4 | def __init__(self, destinations=None): 5 | self.logger = logging.getLogger('Monitoring') 6 | 7 | self.destinations = destinations or [] 8 | 9 | def call_destinations(self, fn_name, *args): 10 | for destination in self.destinations: 11 | try: 12 | getattr(destination, fn_name)(*args) 13 | except: 14 | self.logger.exception('Exception trying to call `{}` on the `{}` monitor.'\ 15 | .format(fn_name, destination.__class__.__name__)) 16 | 17 | def heartbeat_scheduler(self, *args): 18 | self.call_destinations('heartbeat_scheduler', *args) 19 | 20 | def task_retry(self, *args): 21 | self.call_destinations('task_retry', *args) 22 | 23 | def task_failed(self, *args): 24 | self.call_destinations('task_failed', *args) 25 | 26 | def task_success(self, *args): 27 | self.call_destinations('task_success', *args) 28 | 29 | def workflow_failed(self, *args): 30 | self.call_destinations('workflow_failed', *args) 31 | 32 | def workflow_success(self, *args): 33 | self.call_destinations('workflow_success', *args) 34 | 35 | 36 | class MonitorDestination(object): 37 | def heartbeat_scheduler(self, session): 38 | pass 39 | 40 | def task_retry(self, session, task_instance): 41 | pass 42 | 43 | def task_failed(self, session, task_instance): 44 | pass 45 | 46 | def task_success(self, session, task_instance): 47 | pass 48 | 49 | def workflow_failed(self, session, workflow_instance): 50 | pass 51 | 52 | def workflow_success(self, session, workflow_instance): 53 | pass 54 | -------------------------------------------------------------------------------- /taskflow/monitoring/slack.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | import json 4 | 5 | from cron_descriptor import get_description 6 | import requests 7 | 8 | from taskflow import WorkflowInstance, TaskInstance 9 | from .base import MonitorDestination 10 | 11 | class SlackMonitor(MonitorDestination): 12 | def __init__(self, taskflow, slack_url=None): 13 | self.slack_url = slack_url or os.getenv('SLACK_WEBHOOK_URL') 14 | self.taskflow = taskflow 15 | 16 | def get_log_url(self, failed_task): 17 | push_destination = self.taskflow.get_task(failed_task.task_name).push_destination 18 | push_worker = self.taskflow.get_push_worker(push_destination) 19 | if push_worker and hasattr(push_worker, 'get_log_url'): 20 | return push_worker.get_log_url(failed_task) 21 | return None 22 | 23 | def get_message(self, session, item): 24 | failed_task_attachments = None 25 | 26 | if isinstance(item, WorkflowInstance): 27 | item_type = 'Workflow' 28 | name = item.workflow_name 29 | 30 | failed_tasks = session.query(TaskInstance)\ 31 | .filter(TaskInstance.workflow_instance_id == item.id, 32 | TaskInstance.status == 'failed')\ 33 | .all() 34 | 35 | failed_task_attachments = [] 36 | for failed_task in failed_tasks: 37 | failed_task_attachments.append({ 38 | 'title': 'Failed Workflow Task', 39 | 'color': '#ff0000', 40 | 'fields': [ 41 | { 42 | 'title': 'Task', 43 | 'value': failed_task.task_name 44 | }, 45 | { 46 | 'title': 'ID', 47 | 'value': failed_task.id 48 | }, 49 | { 50 | 'title': 'Number of Attempts', 51 | 'value': failed_task.attempts 52 | }, 53 | { 54 | 'title': 'Logs', 55 | 'value': self.get_log_url(failed_task) 56 | } 57 | ] 58 | }) 59 | else: 60 | item_type = 'Task' 61 | name = item.task_name 62 | 63 | attachments = [{ 64 | 'title': '{} Failure'.format(item_type), 65 | 'text': ' A {} in Taskflow failed'.format(item_type.lower()), 66 | 'color': '#ff0000', 67 | 'fields': [ 68 | { 69 | 'title': item_type, 70 | 'value': name, 71 | 'short': False 72 | }, 73 | { 74 | 'title': 'ID', 75 | 'value': item.id 76 | }, 77 | { 78 | 'title': 'Priority', 79 | 'value': item.priority 80 | }, 81 | { 82 | 'title': 'Scheduled Run Time', 83 | 'value': '{:%Y-%m-%d %H:%M:%S}'.format(item.run_at) 84 | }, 85 | { 86 | 'title': 'Start Time', 87 | 'value': '{:%Y-%m-%d %H:%M:%S}'.format(item.started_at) 88 | }, 89 | { 90 | 'title': 'Failure Time', 91 | 'value': '{:%Y-%m-%d %H:%M:%S}'.format(item.ended_at) 92 | } 93 | ] 94 | }] 95 | 96 | if item.scheduled: 97 | if isinstance(item, WorkflowInstance): 98 | schedulable = self.taskflow.get_workflow(item.workflow_name) 99 | else: 100 | schedulable = self.taskflow.get_task(item.task_name) 101 | attachments[0]['fields'].append({ 102 | 'title': 'Schedule', 103 | 'value': '{} ({})'.format(get_description(schedulable.schedule), schedulable.schedule) 104 | }) 105 | 106 | if failed_task_attachments: 107 | attachments += failed_task_attachments 108 | else: 109 | attachments[0]['fields'].append({ 110 | 'title': 'Number of Attempts', 111 | 'value': item.attempts 112 | }) 113 | attachments[0]['fields'].append({ 114 | 'title': 'Logs', 115 | 'value': self.get_log_url(item) 116 | }) 117 | 118 | return {'attachments': attachments} 119 | 120 | def send_to_slack(self, message): 121 | requests.post( 122 | self.slack_url, 123 | data=json.dumps(message), 124 | headers={'Content-Type': 'application/json'}) 125 | 126 | def heartbeat_scheduler(self, session): 127 | pass 128 | 129 | def task_retry(self, session, task_instance): 130 | pass 131 | 132 | def task_failed(self, session, task_instance): 133 | ## only alert on tasks not associated with a workflow. 134 | ## Task failures will bubble up to the workflow 135 | if task_instance.workflow_instance_id == None: 136 | message = self.get_message(session, task_instance) 137 | self.send_to_slack(message) 138 | 139 | def task_success(self, session, task_instance): 140 | pass 141 | 142 | def workflow_failed(self, session, workflow_instance): 143 | message = self.get_message(session, workflow_instance) 144 | self.send_to_slack(message) 145 | 146 | def workflow_success(self, session, workflow_instance): 147 | pass 148 | -------------------------------------------------------------------------------- /taskflow/push_workers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CityOfPhiladelphia/taskflow/5c86b50c6d8bce43a04fba1b856b1d56f3c36b6a/taskflow/push_workers/__init__.py -------------------------------------------------------------------------------- /taskflow/push_workers/aws_batch.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import boto3 4 | 5 | from .base import PushWorker 6 | 7 | class AWSBatchPushWorker(PushWorker): 8 | supports_state_sync = True 9 | push_type = 'aws_batch' 10 | 11 | def __init__(self, *args, default_job_queue=None, default_job_definition=None, **kwargs): 12 | self.logger = logging.getLogger('AWSBatchPushWorker') 13 | 14 | super(AWSBatchPushWorker, self).__init__(*args, **kwargs) 15 | 16 | self.batch_client = boto3.client('batch') 17 | self.default_job_queue = default_job_queue 18 | self.default_job_definition = default_job_definition 19 | 20 | def get_log_url(self, task_instance): 21 | try: 22 | push_state = task_instance.push_state 23 | if push_state and 'taskArn' in push_state: 24 | task_id = re.match(r'task/(.+)', push_state['taskArn']).groups()[0] 25 | 26 | return '{}/{}/{}'.format(push_state['jobName'][:50], push_state['jobId'], task_id) 27 | except: 28 | self.logger.exception('Exception getting AWS Batch log URL') 29 | 30 | return None 31 | 32 | def sync_task_instance_states(self, session, dry_run, task_instances, now): 33 | jobs = dict() 34 | for task_instance in task_instances: 35 | jobs[task_instance.push_state['jobId']] = task_instance 36 | 37 | response = self.batch_client.describe_jobs(jobs=list(jobs.keys())) ## TODO: batch by 100 38 | 39 | ## TODO: tasks that go past 24 hour period ? 40 | ## TODO: timeout pushed tasks? 41 | ## TODO: tasks that are 'pushed' but not in AWS batch? 42 | 43 | for job in response['jobs']: 44 | if job['status'] in ['SUBMITTED','PENDING','RUNNABLE']: 45 | status = 'pushed' 46 | elif job['status'] in ['STARTING','RUNNING']: 47 | status = 'running' 48 | elif job['status'] == 'SUCCEEDED': 49 | status = 'success' 50 | elif job['status'] == 'FAILED': 51 | status = 'failed' 52 | 53 | task_instance = jobs[job['jobId']] 54 | if task_instance.status != status: 55 | task_instance.push_state = job 56 | 57 | if status == 'failed': 58 | task_instance.fail(session, self.taskflow, now=now, dry_run=dry_run) 59 | else: 60 | task_instance.status = status 61 | 62 | if not dry_run: 63 | session.commit() 64 | 65 | def get_job_name(self, workflow, task, task_instance): 66 | if workflow != None: 67 | return '{}__{}__{}__{}'.format( 68 | workflow.name, 69 | task_instance.workflow_instance_id, 70 | task.name, 71 | task_instance.id) 72 | else: 73 | return '{}__{}'.format( 74 | task.name, 75 | task_instance.id) 76 | 77 | def push_task_instances(self, session, dry_run, task_instances): 78 | for task_instance in task_instances: 79 | try: 80 | task = self.taskflow.get_task(task_instance.task_name) 81 | workflow = None 82 | 83 | if task.workflow != None: 84 | workflow = self.taskflow.get_workflow(task.workflow.name) 85 | 86 | if task == None: 87 | raise Exception('Task `{}` not found'.format(task_instance.task_name)) 88 | 89 | parameters = { 90 | 'task': task.name, 91 | 'task_instance': str(task_instance.id) 92 | } 93 | 94 | if workflow != None: 95 | parameters['workflow'] = workflow.name 96 | parameters['workflow_instance'] = str(task_instance.workflow_instance_id) 97 | 98 | if 'job_queue' in task_instance.params and task_instance.params['job_queue']: 99 | job_queue = task_instance.params['job_queue'] 100 | elif 'job_queue' in task.params and task.params['job_queue']: 101 | job_queue = task.params['job_queue'] 102 | else: 103 | job_queue = self.default_job_queue 104 | 105 | if 'job_definition' in task_instance.params and task_instance.params['job_definition']: 106 | job_definition = task_instance.params['job_definition'] 107 | elif 'job_definition' in task.params and task.params['job_definition']: 108 | job_definition = task.params['job_definition'] 109 | else: 110 | job_definition = self.default_job_definition 111 | 112 | environment = [ 113 | { 114 | 'name': 'TASKFLOW_TASK', 115 | 'value': task.name 116 | }, 117 | { 118 | 'name': 'TASKFLOW_TASK_INSTANCE_ID', 119 | 'value': str(task_instance.id) 120 | } 121 | ] 122 | 123 | if workflow != None: 124 | environment.append({ 125 | 'name': 'TASKFLOW_WORKFLOW', 126 | 'value': workflow.name 127 | }) 128 | environment.append({ 129 | 'name': 'TASKFLOW_WORKFLOW_INSTANCE_ID', 130 | 'value': str(task_instance.workflow_instance_id) 131 | }) 132 | 133 | job_name = self.get_job_name(workflow, task, task_instance) 134 | job_name = job_name[-128:] ## AWS Batch max job name is 128 characters 135 | 136 | self.logger.info('Submitting job: %s %s %s', job_name, job_queue, job_definition) 137 | 138 | if not dry_run: 139 | response = self.batch_client.submit_job( 140 | jobName=job_name, 141 | jobQueue=job_queue, 142 | jobDefinition=job_definition, 143 | parameters=parameters, 144 | containerOverrides={ 145 | 'environment': environment 146 | }) 147 | 148 | task_instance.state = 'pushed' 149 | task_instance.push_state = response 150 | 151 | if not dry_run: 152 | session.commit() 153 | except Exception: 154 | self.logger.exception('Exception submitting %s %s', task_instance.task_name, task_instance.id) 155 | session.rollback() 156 | -------------------------------------------------------------------------------- /taskflow/push_workers/base.py: -------------------------------------------------------------------------------- 1 | 2 | class PushWorker(object): 3 | supports_state_sync = False 4 | push_type = None 5 | 6 | def __init__(self, taskflow): 7 | self.taskflow = taskflow 8 | 9 | def sync_task_instance_states(self, session, task_instances): 10 | raise NotImplementedError() 11 | 12 | def push_task_instances(self, session, task_instances): 13 | raise NotImplementedError() 14 | -------------------------------------------------------------------------------- /taskflow/rest/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CityOfPhiladelphia/taskflow/5c86b50c6d8bce43a04fba1b856b1d56f3c36b6a/taskflow/rest/__init__.py -------------------------------------------------------------------------------- /taskflow/rest/app.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import flask 4 | from flask_restful import Api 5 | from flask_cors import CORS 6 | from flask_sqlalchemy import SQLAlchemy 7 | from flask_login import LoginManager 8 | 9 | from taskflow.rest import resources 10 | from taskflow.core.models import metadata, BaseModel, User 11 | 12 | def create_app(taskflow_instance, connection_string=None, secret_key=None): 13 | app = flask.Flask(__name__) 14 | app.config['DEBUG'] = os.getenv('DEBUG', False) 15 | app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False 16 | app.config['SQLALCHEMY_DATABASE_URI'] = connection_string or os.getenv('SQL_ALCHEMY_CONNECTION') 17 | app.config['SESSION_COOKIE_NAME'] = 'taskflowsession' 18 | app.config['SESSION_COOKIE_HTTPONLY'] = True 19 | app.config['PERMANENT_SESSION_LIFETIME'] = 43200 20 | app.config['SECRET_KEY'] = secret_key or os.getenv('FLASK_SESSION_SECRET_KEY') 21 | 22 | db = SQLAlchemy(metadata=metadata, model_class=BaseModel) 23 | 24 | db.init_app(app) 25 | api = Api(app) 26 | CORS(app, supports_credentials=True) 27 | 28 | login_manager = LoginManager() 29 | login_manager.init_app(app) 30 | 31 | @login_manager.user_loader 32 | def load_user(user_id): 33 | return db.session.query(User).filter(User.id == user_id).first() 34 | 35 | def apply_attrs(class_def, attrs): 36 | for key, value in attrs.items(): 37 | setattr(class_def, key, value) 38 | return class_def 39 | 40 | attrs = { 41 | 'session': db.session, 42 | 'taskflow': taskflow_instance 43 | } 44 | 45 | with app.app_context(): 46 | api.add_resource(apply_attrs(resources.LocalSessionResource, attrs), '/v1/session') 47 | 48 | api.add_resource(apply_attrs(resources.WorkflowListResource, attrs), '/v1/workflows') 49 | api.add_resource(apply_attrs(resources.WorkflowResource, attrs), '/v1/workflows/') 50 | 51 | api.add_resource(apply_attrs(resources.TaskListResource, attrs), '/v1/tasks') 52 | api.add_resource(apply_attrs(resources.TaskResource, attrs), '/v1/tasks/') 53 | 54 | api.add_resource(apply_attrs(resources.WorkflowInstanceListResource, attrs), '/v1/workflow-instances') 55 | api.add_resource(apply_attrs(resources.WorkflowInstanceResource, attrs), '/v1/workflow-instances/') 56 | 57 | api.add_resource(apply_attrs(resources.RecurringWorkflowLastestResource, attrs), '/v1/workflow-instances/recurring-latest') 58 | 59 | api.add_resource(apply_attrs(resources.TaskInstanceListResource, attrs), '/v1/task-instances') 60 | api.add_resource(apply_attrs(resources.TaskInstanceResource, attrs), '/v1/task-instances/') 61 | 62 | api.add_resource(apply_attrs(resources.RecurringTaskLastestResource, attrs), '/v1/task-instances/recurring-latest') 63 | 64 | return app 65 | -------------------------------------------------------------------------------- /taskflow/rest/resources.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | from marshmallow import Schema, fields 4 | from marshmallow_sqlalchemy import ModelSchema, field_for 5 | from sqlalchemy import text 6 | from flask import request 7 | from flask_restful import Resource, abort 8 | from flask_login import login_required 9 | from restful_ben.resources import ( 10 | RetrieveUpdateDeleteResource, 11 | QueryEngineMixin, 12 | CreateListResource 13 | ) 14 | from restful_ben.auth import ( 15 | SessionResource, 16 | authorization, 17 | CSRF 18 | ) 19 | 20 | from taskflow import Taskflow, Workflow, WorkflowInstance, Task, TaskInstance 21 | from taskflow.core.models import User 22 | 23 | csrf = CSRF() 24 | 25 | class LocalSessionResource(SessionResource): 26 | User = User 27 | csrf = csrf 28 | 29 | standard_authorization = authorization({ 30 | 'normal': ['GET'], 31 | 'admin': ['POST','PUT','GET','DELETE'] 32 | }) 33 | 34 | def to_list_response(data): 35 | count = len(data) 36 | return { 37 | 'data': data, 38 | 'count': count, 39 | 'total_pages': 1 if count > 0 else 0, 40 | 'page': 1 if count > 0 else 0, 41 | } 42 | 43 | class SchedulableSchema(Schema): 44 | name = fields.String(dump_only=True) 45 | active = fields.Boolean(required=True) 46 | title = fields.String(dump_only=True) 47 | description = fields.String(dump_only=True) 48 | concurrency = fields.Integer(dump_only=True) 49 | sla = fields.Integer(dump_only=True) 50 | schedule = fields.String(dump_only=True) 51 | default_priority = fields.String(dump_only=True) 52 | start_date = fields.DateTime(dump_only=True) 53 | end_date = fields.DateTime(dump_only=True) 54 | 55 | class TaskSchema(SchedulableSchema): 56 | workflow_name = fields.String(dump_only=True) 57 | 58 | class WorkflowSchema(SchedulableSchema): 59 | pass 60 | 61 | workflow_schema = WorkflowSchema() 62 | workflows_schema = WorkflowSchema(many=True) 63 | 64 | class WorkflowListResource(Resource): 65 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 66 | 67 | def get(self): 68 | self.taskflow.sync_db(self.session, read_only=True) 69 | workflows = sorted(self.taskflow.get_workflows(), key=lambda workflow: workflow.name) 70 | workflows_data = workflows_schema.dump(workflows).data 71 | return to_list_response(workflows_data) 72 | 73 | class WorkflowResource(Resource): 74 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 75 | 76 | def get(self, workflow_name): 77 | self.taskflow.sync_db(self.session, read_only=True) 78 | workflow = self.taskflow.get_workflow(workflow_name) 79 | return workflow_schema.dump(workflow).data 80 | 81 | def put(self, workflow_name): 82 | input_workflow = workflow_schema.load(request.json or {}) 83 | 84 | if input_workflow.errors: 85 | abort(400, errors=input_workflow.errors) 86 | 87 | self.taskflow.sync_db(self.session, read_only=True) 88 | workflow = self.taskflow.get_workflow(workflow_name) 89 | 90 | if not workflow: 91 | abort(404) 92 | 93 | workflow.active = input_workflow.data['active'] 94 | self.session.commit() 95 | 96 | return workflow_schema.dump(workflow).data 97 | 98 | task_schema = TaskSchema() 99 | tasks_schema = TaskSchema(many=True) 100 | 101 | class TaskListResource(Resource): 102 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 103 | 104 | def get(self): 105 | self.taskflow.sync_db(self.session, read_only=True) 106 | tasks = list(self.taskflow.get_tasks()) 107 | for workflow in self.taskflow.get_workflows(): 108 | tasks += workflow.get_tasks() 109 | tasks_sorted = sorted(tasks, key=lambda task: task.name) 110 | tasks_data = tasks_schema.dump(tasks_sorted).data 111 | return to_list_response(tasks_data) 112 | 113 | class TaskResource(Resource): 114 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 115 | 116 | def get(self, task_name): 117 | self.taskflow.sync_db(self.session, read_only=True) 118 | task = self.taskflow.get_task(task_name) 119 | return task_schema.dump(task).data 120 | 121 | def put(self, task_name): 122 | input_task = task_schema.load(request.json or {}) 123 | 124 | if input_task.errors: 125 | abort(400, errors=input_task.errors) 126 | 127 | self.taskflow.sync_db(self.session, read_only=True) 128 | task = self.taskflow.get_task(task_name) 129 | 130 | if not task: 131 | abort(404) 132 | 133 | task.active = input_task.data['active'] 134 | self.session.commit() 135 | 136 | return task_schema.dump(task).data 137 | 138 | ## Instances 139 | 140 | class WorkflowInstanceSchema(ModelSchema): 141 | class Meta: 142 | model = WorkflowInstance 143 | exclude = ['task_instances'] 144 | 145 | id = field_for(WorkflowInstance, 'id', dump_only=True) 146 | status = field_for(WorkflowInstance, 'status', dump_only=True) 147 | scheduled = field_for(WorkflowInstance, 'scheduled', dump_only=True) 148 | created_at = field_for(WorkflowInstance, 'created_at', dump_only=True) 149 | updated_at = field_for(WorkflowInstance, 'updated_at', dump_only=True) 150 | 151 | workflow_instance_schema = WorkflowInstanceSchema() 152 | workflow_instances_schema = WorkflowInstanceSchema(many=True) 153 | 154 | class WorkflowInstanceResource(RetrieveUpdateDeleteResource): 155 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 156 | single_schema = workflow_instance_schema 157 | model = WorkflowInstance 158 | 159 | class WorkflowInstanceListResource(QueryEngineMixin, CreateListResource): 160 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 161 | single_schema = workflow_instance_schema 162 | many_schema = workflow_instances_schema 163 | model = WorkflowInstance 164 | 165 | def post(self): 166 | if 'workflow_name' in request.json: 167 | workflow = self.taskflow.get_workflow(request.json['workflow_name']) 168 | if not workflow: 169 | abort(400, errors={'workflow_name': 'Workflow `{}` not found'.format(request.json['workflow_name'])}) 170 | 171 | if 'priority' not in request.json: 172 | request.json['priority'] = workflow.default_priority 173 | 174 | return super(WorkflowInstanceListResource, self).post() 175 | 176 | class RecurringWorkflowLastestResource(Resource): 177 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 178 | 179 | def get(self): 180 | workflow_instances = self.session.query(WorkflowInstance)\ 181 | .from_statement(text(""" 182 | SELECT workflow_instances.* FROM workflow_instances 183 | INNER JOIN ( 184 | SELECT workflow_name, MAX(started_at) AS max_date 185 | FROM workflow_instances 186 | WHERE status != 'queued' AND scheduled = true 187 | GROUP BY workflow_name) AS tp 188 | ON workflow_instances.workflow_name = tp.workflow_name AND 189 | workflow_instances.started_at = tp.max_date 190 | ORDER BY workflow_instances.workflow_name; 191 | """))\ 192 | .all() 193 | 194 | return workflow_instances_schema.dump(workflow_instances).data 195 | 196 | class TaskInstanceSchema(ModelSchema): 197 | class Meta: 198 | model = TaskInstance 199 | exclude = ['workflow'] 200 | 201 | id = field_for(TaskInstance, 'id', dump_only=True) 202 | workflow_instance_id = field_for(TaskInstance, 'workflow_instance_id', dump_only=True) 203 | status = field_for(TaskInstance, 'status', dump_only=True) 204 | scheduled = field_for(WorkflowInstance, 'scheduled', dump_only=True) 205 | created_at = field_for(TaskInstance, 'created_at', dump_only=True) 206 | updated_at = field_for(TaskInstance, 'updated_at', dump_only=True) 207 | 208 | task_instance_schema = TaskInstanceSchema() 209 | task_instances_schema = TaskInstanceSchema(many=True) 210 | 211 | class TaskInstanceResource(RetrieveUpdateDeleteResource): 212 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 213 | single_schema = task_instance_schema 214 | model = TaskInstance 215 | 216 | class TaskInstanceListResource(QueryEngineMixin, CreateListResource): 217 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 218 | single_schema = task_instance_schema 219 | many_schema = task_instances_schema 220 | model = TaskInstance 221 | 222 | def post(self): 223 | if 'task_name' in request.json: 224 | task = self.taskflow.get_task(request.json['task_name']) 225 | if not task: 226 | abort(400, errors={'task_name': 'Task `{}` not found'.format(request.json['task_name'])}) 227 | 228 | request.json['push'] = task.push_destination != None 229 | 230 | if 'priority' not in request.json: 231 | request.json['priority'] = task.default_priority 232 | 233 | if 'max_attempts' not in request.json: 234 | request.json['max_attempts'] = task.retries + 1 235 | 236 | if 'timeout' not in request.json: 237 | request.json['timeout'] = task.timeout 238 | 239 | if 'retry_delay' not in request.json: 240 | request.json['retry_delay'] = task.retry_delay 241 | 242 | return super(TaskInstanceListResource, self).post() 243 | 244 | class RecurringTaskLastestResource(Resource): 245 | method_decorators = [csrf.csrf_check, standard_authorization, login_required] 246 | 247 | def get(self): 248 | task_instances = self.session.query(TaskInstance)\ 249 | .from_statement(text(""" 250 | SELECT task_instances.* FROM task_instances 251 | INNER JOIN ( 252 | SELECT task_name, MAX(started_at) AS max_date 253 | FROM task_instances 254 | WHERE status != 'queued' AND scheduled = true 255 | GROUP BY task_name) AS tp 256 | ON task_instances.task_name = tp.task_name AND 257 | task_instances.started_at = tp.max_date 258 | ORDER BY task_instances.task_name; 259 | """))\ 260 | .all() 261 | 262 | return task_instances_schema.dump(task_instances).data 263 | -------------------------------------------------------------------------------- /taskflow/tasks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CityOfPhiladelphia/taskflow/5c86b50c6d8bce43a04fba1b856b1d56f3c36b6a/taskflow/tasks/__init__.py -------------------------------------------------------------------------------- /taskflow/tasks/bash_task.py: -------------------------------------------------------------------------------- 1 | import os 2 | import signal 3 | import logging 4 | import re 5 | import sys 6 | import shutil 7 | import time 8 | from subprocess import Popen, PIPE 9 | from tempfile import gettempdir, NamedTemporaryFile 10 | from threading import Thread 11 | from contextlib import contextmanager 12 | from tempfile import mkdtemp 13 | 14 | from smart_open import smart_open 15 | import boto 16 | import boto3 17 | 18 | from taskflow import Task 19 | 20 | def pipe_stream(stream1, stream2): 21 | def stream_helper(stream1, stream2): 22 | for line in iter(stream1.readline, b''): 23 | stream2.write(line) 24 | 25 | t = Thread(target=stream_helper, args=(stream1, stream2)) 26 | t.daemon = True 27 | t.start() 28 | 29 | return t 30 | 31 | @contextmanager 32 | def TemporaryDirectory(suffix='', prefix=None, dir=None): 33 | name = mkdtemp(suffix=suffix, prefix=prefix, dir=dir) 34 | try: 35 | yield name 36 | finally: 37 | try: 38 | shutil.rmtree(name) 39 | except OSError as e: 40 | # ENOENT - no such file or directory 41 | if e.errno != errno.ENOENT: 42 | raise e 43 | 44 | def replace_environment_variables(input_str): 45 | env_vars = re.findall(r'"?\$([A-Z0-9_-]+)"?', input_str) 46 | 47 | out_str = input_str 48 | for env_var in env_vars: 49 | value = os.getenv(env_var, '') 50 | out_str = re.sub('"?\$' + env_var + '"?', value, out_str) 51 | return out_str 52 | 53 | class BashTask(Task): 54 | def get_command(self): 55 | return self.params['command'] 56 | 57 | def execute(self, task_instance): 58 | logger = logging.getLogger(self.name) 59 | bash_command = self.get_command() 60 | logger.info('Temporary directory root location: %s', gettempdir()) 61 | with TemporaryDirectory(prefix='taskflowtmp') as tmp_dir: 62 | with NamedTemporaryFile(dir=tmp_dir, prefix=str(task_instance.id)) as f: 63 | f.write(bytes(bash_command, 'utf_8')) 64 | f.flush() 65 | fname = f.name 66 | script_location = tmp_dir + "/" + fname 67 | logger.info('Temporary script location: %s', script_location) 68 | logger.info('Running command: %s', bash_command) 69 | 70 | inputpath = None 71 | if 'input_file' in task_instance.params and task_instance.params['input_file'] != None: 72 | inputpath = replace_environment_variables(task_instance.params['input_file']) 73 | elif 'input_file' in self.params and self.params['input_file'] != None: 74 | inputpath = replace_environment_variables(self.params['input_file']) 75 | 76 | input_file = None 77 | if inputpath: 78 | logger.info('Streaming to STDIN from: %s', inputpath) 79 | input_file = smart_open(inputpath, mode='rb') 80 | 81 | outpath = None 82 | if 'output_file' in task_instance.params and task_instance.params['output_file'] != None: 83 | outpath = replace_environment_variables(task_instance.params['output_file']) 84 | elif 'output_file' in self.params and self.params['output_file'] != None: 85 | outpath = replace_environment_variables(self.params['output_file']) 86 | 87 | output_file = None 88 | if outpath: 89 | logger.info('Streaming STDOUT to: %s', outpath) 90 | output_file = smart_open(outpath, mode='wb') 91 | 92 | ON_POSIX = 'posix' in sys.builtin_module_names 93 | 94 | sp = Popen( 95 | ['bash', fname], 96 | stdin=PIPE if input_file else None, 97 | stdout=PIPE if output_file else None, 98 | stderr=PIPE, 99 | cwd=tmp_dir, 100 | preexec_fn=os.setsid, 101 | bufsize=1, 102 | close_fds=ON_POSIX) 103 | 104 | self.sp = sp 105 | 106 | input_thread = None 107 | if input_file: 108 | input_thread = pipe_stream(input_file, sp.stdin) 109 | 110 | output_thread = None 111 | if output_file: 112 | output_thread = pipe_stream(sp.stdout, output_file) 113 | 114 | for line in iter(sp.stderr.readline, b''): 115 | logger.info(line) 116 | 117 | sp.wait() 118 | 119 | if input_thread: 120 | input_thread.join(timeout=5) 121 | 122 | if output_thread: 123 | output_thread.join(timeout=5) 124 | 125 | if input_file: 126 | input_file.close() 127 | 128 | if output_file: 129 | logger.info('Closing STDOUT file') 130 | start = time.time() 131 | output_file.close() 132 | logger.info('STDOUT file written - {}s'.format(time.time() - start)) 133 | 134 | logger.info('Command exited with return code %s', sp.returncode) 135 | 136 | if sp.returncode: 137 | raise Exception('Bash command failed') 138 | 139 | def on_kill(self): 140 | logging.info('Sending SIGTERM signal to bash process group') 141 | os.killpg(os.getpgid(self.sp.pid), signal.SIGTERM) 142 | -------------------------------------------------------------------------------- /tests/shared_fixtures.py: -------------------------------------------------------------------------------- 1 | import os 2 | import logging 3 | from datetime import datetime 4 | 5 | from sqlalchemy import create_engine 6 | from sqlalchemy.orm import Session 7 | import pytest 8 | 9 | os.environ['CSRF_SECRET'] = 'test' 10 | 11 | from taskflow import Taskflow, Workflow, WorkflowInstance, Task, TaskInstance 12 | from taskflow.core.models import BaseModel, User 13 | from taskflow.rest.app import create_app 14 | 15 | def get_logging(): 16 | logger = logging.getLogger() 17 | handler = logging.StreamHandler() 18 | formatter = logging.Formatter('[%(asctime)s] %(name)s %(levelname)s %(message)s') 19 | handler.setFormatter(formatter) 20 | logger.addHandler(handler) 21 | logger.setLevel(logging.DEBUG) 22 | 23 | @pytest.fixture(scope='session') 24 | def engine(): 25 | return create_engine('postgresql://localhost/taskflow_test') 26 | 27 | @pytest.fixture 28 | def tables(engine): 29 | BaseModel.metadata.create_all(engine) 30 | yield 31 | BaseModel.metadata.drop_all(engine) 32 | 33 | @pytest.fixture 34 | def dbsession(engine, tables): 35 | """Returns an sqlalchemy session, and after the test tears down everything properly.""" 36 | connection = engine.connect() 37 | # use the connection with the already started transaction 38 | session = Session(bind=connection) 39 | 40 | yield session 41 | 42 | session.close() 43 | # put back the connection to the connection pool 44 | connection.close() 45 | 46 | @pytest.fixture 47 | def workflows(dbsession): 48 | workflow1 = Workflow(name='workflow1', active=True, schedule='0 6 * * *') 49 | workflow2 = Workflow(name='workflow2', active=True) 50 | dbsession.add(workflow1) 51 | dbsession.add(workflow2) 52 | dbsession.commit() 53 | 54 | task1 = Task(workflow=workflow1, name='task1', active=True) 55 | task2 = Task(workflow=workflow1, name='task2', active=True) 56 | task3 = Task(workflow=workflow1, name='task3', active=True) 57 | task4 = Task(workflow=workflow1, name='task4', active=True) 58 | 59 | task3.depends_on(task1) 60 | task3.depends_on(task2) 61 | task4.depends_on(task3) 62 | 63 | dbsession.add(task1) 64 | dbsession.add(task2) 65 | dbsession.add(task3) 66 | dbsession.add(task4) 67 | 68 | dbsession.commit() 69 | return [workflow1, workflow2] 70 | 71 | @pytest.fixture 72 | def instances(dbsession, workflows): 73 | workflow_instance = WorkflowInstance( 74 | workflow_name='workflow1', 75 | scheduled=True, 76 | run_at=datetime(2017, 6, 3, 6), 77 | started_at=datetime(2017, 6, 3, 6), 78 | status='running', 79 | priority='normal') 80 | dbsession.add(workflow_instance) 81 | dbsession.commit() 82 | task_instance1 = TaskInstance( 83 | task_name='task1', 84 | scheduled=True, 85 | workflow_instance_id=workflow_instance.id, 86 | status='success', 87 | run_at=datetime(2017, 6, 3, 6, 0, 12), 88 | started_at=datetime(2017, 6, 3, 6, 0, 12), 89 | ended_at=datetime(2017, 6, 3, 6, 0, 18), 90 | attempts=1, 91 | priority='normal', 92 | push=False, 93 | timeout=300, 94 | retry_delay=300) 95 | task_instance2 = TaskInstance( 96 | task_name='task2', 97 | scheduled=True, 98 | workflow_instance_id=workflow_instance.id, 99 | status='success', 100 | run_at=datetime(2017, 6, 3, 6, 0, 20), 101 | started_at=datetime(2017, 6, 3, 6, 0, 20), 102 | ended_at=datetime(2017, 6, 3, 6, 0, 27), 103 | attempts=1, 104 | priority='normal', 105 | push=False, 106 | timeout=300, 107 | retry_delay=300) 108 | task_instance3 = TaskInstance( 109 | task_name='task3', 110 | scheduled=True, 111 | workflow_instance_id=workflow_instance.id, 112 | status='success', 113 | run_at=datetime(2017, 6, 3, 6, 0, 28), 114 | started_at=datetime(2017, 6, 3, 6, 0, 28), 115 | ended_at=datetime(2017, 6, 3, 6, 0, 32), 116 | attempts=1, 117 | priority='normal', 118 | push=False, 119 | timeout=300, 120 | retry_delay=300) 121 | task_instance4 = TaskInstance( 122 | task_name='task4', 123 | scheduled=True, 124 | workflow_instance_id=workflow_instance.id, 125 | status='running', 126 | run_at=datetime(2017, 6, 3, 6, 0, 34), 127 | started_at=datetime(2017, 6, 3, 6, 0, 34), 128 | attempts=1, 129 | priority='normal', 130 | push=False, 131 | timeout=300, 132 | retry_delay=300) 133 | dbsession.add(task_instance1) 134 | dbsession.add(task_instance2) 135 | dbsession.add(task_instance3) 136 | dbsession.add(task_instance4) 137 | dbsession.commit() 138 | 139 | @pytest.fixture 140 | def app(tables, workflows, dbsession): 141 | taskflow = Taskflow() 142 | taskflow.add_workflows(workflows) 143 | taskflow.sync_db(dbsession) 144 | 145 | dbsession.add(User(active=True, username='amadonna', password='foo', role='admin')) 146 | dbsession.commit() 147 | 148 | app = create_app(taskflow, connection_string='postgresql://localhost/taskflow_test', secret_key='foo') 149 | 150 | yield app 151 | -------------------------------------------------------------------------------- /tests/test_basic_workflow.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | import pytest 4 | from sqlalchemy.exc import IntegrityError 5 | 6 | from taskflow import Scheduler, Taskflow, Workflow, WorkflowInstance, Task, TaskInstance 7 | from shared_fixtures import * 8 | 9 | get_logging() 10 | 11 | ## TODO: test dry run 12 | 13 | def test_schedule_recurring_workflow(dbsession, workflows): 14 | taskflow = Taskflow() 15 | taskflow.add_workflows(workflows) 16 | scheduler = Scheduler(taskflow, now_override=datetime(2017, 6, 3, 6)) 17 | scheduler.run(dbsession) 18 | 19 | workflow_instances = dbsession.query(WorkflowInstance).all() 20 | assert len(workflow_instances) == 1 21 | assert workflow_instances[0].status == 'queued' 22 | assert workflow_instances[0].scheduled == True 23 | assert workflow_instances[0].run_at == datetime(2017, 6, 4, 6) 24 | 25 | task_instances = dbsession.query(TaskInstance).all() 26 | assert len(task_instances) == 0 27 | 28 | def test_unique_workflow(dbsession): 29 | workflow1 = Workflow(name='workflow1', active=True) 30 | dbsession.add(workflow1) 31 | dbsession.commit() 32 | 33 | workflow_instance = WorkflowInstance( 34 | workflow_name='workflow1', 35 | scheduled=False, 36 | run_at=datetime(2017, 6, 3, 6), 37 | status='queued', 38 | priority='normal', 39 | unique='foo') 40 | dbsession.add(workflow_instance) 41 | dbsession.commit() 42 | 43 | with pytest.raises(IntegrityError): 44 | workflow_instance = WorkflowInstance( 45 | workflow_name='workflow1', 46 | scheduled=False, 47 | run_at=datetime(2017, 6, 3, 6), 48 | status='queued', 49 | priority='normal', 50 | unique='foo') 51 | dbsession.add(workflow_instance) 52 | dbsession.commit() 53 | 54 | def test_workflow_starts(dbsession, workflows): 55 | now = datetime(2017, 6, 3, 6, 12) 56 | 57 | taskflow = Taskflow() 58 | taskflow.add_workflows(workflows) 59 | 60 | workflow_instance = WorkflowInstance( 61 | workflow_name='workflow1', 62 | scheduled=True, 63 | run_at=datetime(2017, 6, 3, 6), 64 | status='queued', 65 | priority='normal') 66 | dbsession.add(workflow_instance) 67 | dbsession.commit() 68 | 69 | scheduler = Scheduler(taskflow, now_override=now) 70 | scheduler.run(dbsession) 71 | 72 | assert workflow_instance.status == 'running' 73 | assert workflow_instance.started_at == now 74 | 75 | task_instances = dbsession.query(TaskInstance).all() 76 | assert len(task_instances) == 2 77 | for instance in task_instances: 78 | assert instance.task_name in ['task1','task2'] 79 | assert instance.status == 'queued' 80 | 81 | def test_schedule_recurring_workflow(dbsession, workflows): 82 | taskflow = Taskflow() 83 | taskflow.add_workflows(workflows) 84 | 85 | workflow_instance = WorkflowInstance( 86 | workflow_name='workflow1', 87 | scheduled=True, 88 | run_at=datetime(2017, 6, 3, 6), 89 | status='queued', 90 | priority='normal') 91 | dbsession.add(workflow_instance) 92 | dbsession.commit() 93 | 94 | scheduler = Scheduler(taskflow, now_override=datetime(2017, 6, 3, 6, 0, 45)) 95 | scheduler.run(dbsession) 96 | 97 | dbsession.refresh(workflow_instance) 98 | assert workflow_instance.status == 'running' 99 | 100 | task_instances = dbsession.query(TaskInstance).all() 101 | assert len(task_instances) == 2 102 | for instance in task_instances: 103 | assert instance.task_name in ['task1','task2'] 104 | assert instance.status == 'queued' 105 | 106 | def test_workflow_running_no_change(dbsession, workflows): 107 | taskflow = Taskflow() 108 | taskflow.add_workflows(workflows) 109 | 110 | workflow1 = workflows[0] 111 | 112 | workflow_instance = WorkflowInstance( 113 | workflow_name='workflow1', 114 | scheduled=True, 115 | run_at=datetime(2017, 6, 3, 6), 116 | started_at=datetime(2017, 6, 3, 6), 117 | status='running', 118 | priority='normal') 119 | dbsession.add(workflow_instance) 120 | dbsession.commit() 121 | 122 | task_instance1 = TaskInstance( 123 | task_name='task1', 124 | scheduled=True, 125 | workflow_instance_id=workflow_instance.id, 126 | status='running', 127 | run_at=datetime(2017, 6, 3, 6, 0, 34), 128 | attempts=1, 129 | priority='normal', 130 | push=False, 131 | timeout=300, 132 | retry_delay=300) 133 | task_instance2 = TaskInstance( 134 | task_name='task2', 135 | scheduled=True, 136 | workflow_instance_id=workflow_instance.id, 137 | status='running', 138 | run_at=datetime(2017, 6, 3, 6, 0, 34), 139 | attempts=1, 140 | priority='normal', 141 | push=False, 142 | timeout=300, 143 | retry_delay=300) 144 | dbsession.add(task_instance1) 145 | dbsession.add(task_instance2) 146 | dbsession.commit() 147 | 148 | scheduler = Scheduler(taskflow, now_override=datetime(2017, 6, 3, 6, 12)) 149 | scheduler.run(dbsession) 150 | 151 | task_instances = dbsession.query(TaskInstance).all() 152 | assert len(task_instances) == 2 153 | for instance in task_instances: 154 | assert instance.task_name in ['task1','task2'] 155 | assert instance.status == 'running' 156 | 157 | def test_workflow_next_step(dbsession, workflows): 158 | taskflow = Taskflow() 159 | taskflow.add_workflows(workflows) 160 | 161 | workflow_instance = WorkflowInstance( 162 | workflow_name='workflow1', 163 | scheduled=True, 164 | run_at=datetime(2017, 6, 3, 6), 165 | status='running', 166 | priority='normal') 167 | dbsession.add(workflow_instance) 168 | dbsession.commit() 169 | task_instance1 = TaskInstance( 170 | task_name='task1', 171 | scheduled=True, 172 | workflow_instance_id=workflow_instance.id, 173 | status='success', 174 | run_at=datetime(2017, 6, 3, 6, 0, 34), 175 | attempts=1, 176 | priority='normal', 177 | push=False, 178 | timeout=300, 179 | retry_delay=300) 180 | task_instance2 = TaskInstance( 181 | task_name='task2', 182 | scheduled=True, 183 | workflow_instance_id=workflow_instance.id, 184 | status='success', 185 | run_at=datetime(2017, 6, 3, 6, 0, 34), 186 | attempts=1, 187 | priority='normal', 188 | push=False, 189 | timeout=300, 190 | retry_delay=300) 191 | dbsession.add(task_instance1) 192 | dbsession.add(task_instance2) 193 | dbsession.commit() 194 | 195 | scheduler = Scheduler(taskflow, now_override=datetime(2017, 6, 3, 6, 12)) 196 | scheduler.run(dbsession) 197 | 198 | task_instances = dbsession.query(TaskInstance).all() 199 | assert len(task_instances) == 3 200 | for instance in task_instances: 201 | assert instance.task_name in ['task1','task2','task3'] 202 | if instance.task_name in ['task1','task2']: 203 | assert instance.status == 'success' 204 | elif instance.task_name == 'task3': 205 | assert instance.status == 'queued' 206 | 207 | def test_workflow_success(dbsession, workflows): 208 | taskflow = Taskflow() 209 | taskflow.add_workflows(workflows) 210 | 211 | workflow_instance = WorkflowInstance( 212 | workflow_name='workflow1', 213 | scheduled=True, 214 | run_at=datetime(2017, 6, 3, 6), 215 | status='running', 216 | priority='normal') 217 | dbsession.add(workflow_instance) 218 | dbsession.commit() 219 | task_instance1 = TaskInstance( 220 | task_name='task1', 221 | scheduled=True, 222 | workflow_instance_id=workflow_instance.id, 223 | status='success', 224 | run_at=datetime(2017, 6, 3, 6, 0, 34), 225 | attempts=1, 226 | priority='normal', 227 | push=False, 228 | timeout=300, 229 | retry_delay=300) 230 | task_instance2 = TaskInstance( 231 | task_name='task2', 232 | scheduled=True, 233 | workflow_instance_id=workflow_instance.id, 234 | status='success', 235 | run_at=datetime(2017, 6, 3, 6, 0, 34), 236 | attempts=1, 237 | priority='normal', 238 | push=False, 239 | timeout=300, 240 | retry_delay=300) 241 | task_instance3 = TaskInstance( 242 | task_name='task3', 243 | scheduled=True, 244 | workflow_instance_id=workflow_instance.id, 245 | status='success', 246 | run_at=datetime(2017, 6, 3, 6, 0, 34), 247 | attempts=1, 248 | priority='normal', 249 | push=False, 250 | timeout=300, 251 | retry_delay=300) 252 | task_instance4 = TaskInstance( 253 | task_name='task4', 254 | scheduled=True, 255 | workflow_instance_id=workflow_instance.id, 256 | status='success', 257 | run_at=datetime(2017, 6, 3, 6, 0, 34), 258 | attempts=1, 259 | priority='normal', 260 | push=False, 261 | timeout=300, 262 | retry_delay=300) 263 | dbsession.add(task_instance1) 264 | dbsession.add(task_instance2) 265 | dbsession.add(task_instance3) 266 | dbsession.add(task_instance4) 267 | dbsession.commit() 268 | 269 | now = datetime(2017, 6, 3, 6, 12) 270 | 271 | scheduler = Scheduler(taskflow, now_override=now) 272 | scheduler.run(dbsession) 273 | 274 | dbsession.refresh(workflow_instance) 275 | assert workflow_instance.status == 'success' 276 | assert workflow_instance.ended_at == now 277 | 278 | task_instances = dbsession.query(TaskInstance).all() 279 | assert len(task_instances) == 4 280 | for instance in task_instances: 281 | assert instance.task_name in ['task1','task2','task3','task4'] 282 | assert instance.status == 'success' 283 | 284 | def test_workflow_fail(dbsession, workflows): 285 | taskflow = Taskflow() 286 | taskflow.add_workflows(workflows) 287 | 288 | workflow_instance = WorkflowInstance( 289 | workflow_name='workflow1', 290 | scheduled=True, 291 | run_at=datetime(2017, 6, 3, 6), 292 | status='running', 293 | priority='normal') 294 | dbsession.add(workflow_instance) 295 | dbsession.commit() 296 | task_instance1 = TaskInstance( 297 | task_name='task1', 298 | scheduled=True, 299 | workflow_instance_id=workflow_instance.id, 300 | status='success', 301 | run_at=datetime(2017, 6, 3, 6, 0, 34), 302 | attempts=1, 303 | priority='normal', 304 | push=False, 305 | timeout=300, 306 | retry_delay=300) 307 | task_instance2 = TaskInstance( 308 | task_name='task2', 309 | scheduled=True, 310 | workflow_instance_id=workflow_instance.id, 311 | status='success', 312 | run_at=datetime(2017, 6, 3, 6, 0, 34), 313 | attempts=1, 314 | priority='normal', 315 | push=False, 316 | timeout=300, 317 | retry_delay=300) 318 | task_instance3 = TaskInstance( 319 | task_name='task3', 320 | scheduled=True, 321 | workflow_instance_id=workflow_instance.id, 322 | status='failed', 323 | run_at=datetime(2017, 6, 3, 6, 0, 34), 324 | attempts=1, 325 | priority='normal', 326 | push=False, 327 | timeout=300, 328 | retry_delay=300) 329 | dbsession.add(task_instance1) 330 | dbsession.add(task_instance2) 331 | dbsession.add(task_instance3) 332 | dbsession.commit() 333 | 334 | now = datetime(2017, 6, 3, 6, 12) 335 | 336 | scheduler = Scheduler(taskflow, now_override=now) 337 | scheduler.run(dbsession) 338 | 339 | dbsession.refresh(workflow_instance) 340 | assert workflow_instance.status == 'failed' 341 | assert workflow_instance.ended_at == now 342 | 343 | task_instances = dbsession.query(TaskInstance).all() 344 | assert len(task_instances) == 3 345 | for instance in task_instances: 346 | assert instance.task_name in ['task1','task2','task3'] 347 | if instance.task_name in ['task1','task2']: 348 | assert instance.status == 'success' 349 | elif instance.task_name == 'task3': 350 | assert instance.status == 'failed' 351 | -------------------------------------------------------------------------------- /tests/test_monitoring.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | import requests_mock 4 | from restful_ben.test_utils import dict_contains 5 | 6 | from taskflow import WorkflowInstance, TaskInstance 7 | from taskflow.monitoring.slack import SlackMonitor 8 | 9 | from shared_fixtures import * 10 | 11 | def test_slack_monitor(dbsession, workflows): 12 | taskflow = Taskflow() 13 | taskflow.add_workflows(workflows) 14 | 15 | workflow_instance = WorkflowInstance( 16 | workflow_name='workflow1', 17 | scheduled=True, 18 | run_at=datetime(2017, 6, 3, 6), 19 | started_at=datetime(2017, 6, 3, 6), 20 | ended_at=datetime(2017, 6, 3, 6, 0, 36), 21 | status='running', 22 | priority='normal') 23 | dbsession.add(workflow_instance) 24 | dbsession.commit() 25 | task_instance1 = TaskInstance( 26 | task_name='task1', 27 | scheduled=True, 28 | workflow_instance_id=workflow_instance.id, 29 | status='success', 30 | run_at=datetime(2017, 6, 3, 6, 0, 34), 31 | attempts=1, 32 | priority='normal', 33 | push=False, 34 | timeout=300, 35 | retry_delay=300) 36 | task_instance2 = TaskInstance( 37 | task_name='task2', 38 | scheduled=True, 39 | workflow_instance_id=workflow_instance.id, 40 | status='success', 41 | run_at=datetime(2017, 6, 3, 6, 0, 34), 42 | attempts=1, 43 | priority='normal', 44 | push=False, 45 | timeout=300, 46 | retry_delay=300) 47 | task_instance3 = TaskInstance( 48 | task_name='task3', 49 | scheduled=True, 50 | workflow_instance_id=workflow_instance.id, 51 | status='failed', 52 | run_at=datetime(2017, 6, 3, 6, 0, 34), 53 | attempts=1, 54 | priority='normal', 55 | push=False, 56 | timeout=300, 57 | retry_delay=300) 58 | dbsession.add(task_instance1) 59 | dbsession.add(task_instance2) 60 | dbsession.add(task_instance3) 61 | dbsession.commit() 62 | 63 | slack_url = 'http://fakeurl.com/foo' 64 | 65 | slack_monitor = SlackMonitor(taskflow, slack_url=slack_url) 66 | 67 | with requests_mock.Mocker() as m: 68 | m.post(slack_url) 69 | 70 | slack_monitor.workflow_failed(dbsession, workflow_instance) 71 | 72 | assert m.called 73 | assert m.call_count == 1 74 | 75 | request = m.request_history[0] 76 | data = request.json() 77 | assert dict_contains(data['attachments'][0], { 78 | 'text': ' A workflow in Taskflow failed', 79 | 'title': 'Workflow Failure', 80 | 'color': '#ff0000', 81 | 'fields': [ 82 | { 83 | 'title': 'Workflow', 84 | 'short': False, 85 | 'value': 'workflow1' 86 | }, 87 | { 88 | 'title': 'ID', 89 | 'value': 1 90 | }, 91 | { 92 | 'title': 'Priority', 93 | 'value': 'normal' 94 | }, 95 | { 96 | 'title': 'Scheduled Run Time', 97 | 'value': '2017-06-03 06:00:00' 98 | }, 99 | { 100 | 'title': 'Start Time', 101 | 'value': '2017-06-03 06:00:00' 102 | }, 103 | { 104 | 'title': 'Failure Time', 105 | 'value': '2017-06-03 06:00:36' 106 | }, 107 | { 108 | 'title': 'Schedule', 109 | 'value': 'At 06:00 AM (0 6 * * *)' 110 | } 111 | ] 112 | }) 113 | assert dict_contains(data['attachments'][1], { 114 | 'title': 'Failed Workflow Task', 115 | 'color': '#ff0000', 116 | 'fields': [ 117 | { 118 | 'title': 'Task', 119 | 'value': 'task3' 120 | }, 121 | { 122 | 'title': 'ID', 123 | 'value': 3 124 | }, 125 | { 126 | 'title': 'Number of Attempts', 127 | 'value': 1 128 | }, 129 | { 130 | 'title': 'Logs', 131 | 'value': None 132 | } 133 | ] 134 | }) -------------------------------------------------------------------------------- /tests/test_push_aws_batch.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | from functools import reduce 3 | import uuid 4 | 5 | import pytest 6 | import boto3 7 | 8 | from taskflow import Scheduler, Pusher, Taskflow, Task, TaskInstance 9 | from taskflow.push_workers.aws_batch import AWSBatchPushWorker 10 | from shared_fixtures import * 11 | 12 | get_logging() 13 | 14 | class MockAWSBatch(object): 15 | def __init__(self, jobs, status): 16 | self.jobs = jobs 17 | self.status = status 18 | 19 | def describe_jobs(self, *args, **kwargs): 20 | jobs_with_status = self.jobs.copy() 21 | for job in jobs_with_status: 22 | job['status'] = self.status 23 | return {'jobs': jobs_with_status} 24 | 25 | def remove_job(self, job_id): ## only added for testing 26 | for job in self.jobs: 27 | if job['jobId'] == job_id: 28 | self.jobs.remove(job) 29 | break 30 | 31 | def submit_job(self, *args, **job): 32 | job['jobId'] = str(uuid.uuid4()) 33 | self.jobs.append(job) 34 | return { 35 | 'jobName': job['jobName'], 36 | 'jobId': job['jobId'] 37 | } 38 | 39 | def mockbatch(mock_aws_batch): 40 | def mockclient(aws_resource): 41 | if aws_resource != 'batch': 42 | raise Exception('Only AWS Batch is mocked!') 43 | return mock_aws_batch 44 | return mockclient 45 | 46 | # def test_push(dbsession, monkeypatch): 47 | # mock_aws_batch = MockAWSBatch([], 'SUBMITTED') 48 | # monkeypatch.setattr(boto3, 'client', mockbatch(mock_aws_batch)) 49 | 50 | # task1 = Task(name='task1', active=True, push_destination='aws_batch') 51 | # dbsession.add(task1) 52 | # taskflow = Taskflow() 53 | # taskflow.add_task(task1) 54 | # taskflow.add_push_worker(AWSBatchPushWorker(taskflow)) 55 | # taskflow.sync_db(dbsession) 56 | 57 | # task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 58 | # dbsession.add(task_instance) 59 | # dbsession.commit() 60 | 61 | # pusher = Pusher(taskflow, now_override=datetime(2017, 6, 4, 6)) 62 | # pusher.run(dbsession) 63 | 64 | # task_instances = dbsession.query(TaskInstance).all() 65 | # pushed_task_instance = task_instances[0] 66 | # assert pushed_task_instance.status == 'pushed' 67 | # assert 'jobId' in pushed_task_instance.push_state 68 | # assert pushed_task_instance.push_state['jobName'] == 'task1__1' 69 | 70 | # mock_aws_batch.status = 'PENDING' 71 | # pusher.run(dbsession) 72 | # assert pushed_task_instance.status == 'pushed' 73 | 74 | # mock_aws_batch.status = 'RUNNABLE' 75 | # pusher.run(dbsession) 76 | # assert pushed_task_instance.status == 'pushed' 77 | 78 | # mock_aws_batch.status = 'RUNNING' 79 | # pusher.run(dbsession) 80 | # assert pushed_task_instance.status == 'running' 81 | 82 | # mock_aws_batch.status = 'STARTING' 83 | # pusher.run(dbsession) 84 | # assert pushed_task_instance.status == 'running' 85 | 86 | # mock_aws_batch.status = 'SUCCEEDED' 87 | # pusher.run(dbsession) 88 | # assert pushed_task_instance.status == 'success' 89 | 90 | # def test_fail(dbsession, monkeypatch): 91 | # mock_aws_batch = MockAWSBatch([], 'SUBMITTED') 92 | # monkeypatch.setattr(boto3, 'client', mockbatch(mock_aws_batch)) 93 | 94 | # task1 = Task(name='task1', active=True, push_destination='aws_batch') 95 | # dbsession.add(task1) 96 | # taskflow = Taskflow() 97 | # taskflow.add_task(task1) 98 | # taskflow.add_push_worker(AWSBatchPushWorker(taskflow)) 99 | # taskflow.sync_db(dbsession) 100 | 101 | # task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 102 | # dbsession.add(task_instance) 103 | # dbsession.commit() 104 | 105 | # pusher = Pusher(taskflow, now_override=datetime(2017, 6, 4, 6)) 106 | # pusher.run(dbsession) 107 | 108 | # task_instances = dbsession.query(TaskInstance).all() 109 | # pushed_task_instance = task_instances[0] 110 | # assert pushed_task_instance.status == 'pushed' 111 | # assert 'jobId' in pushed_task_instance.push_state 112 | # assert pushed_task_instance.push_state['jobName'] == 'task1__1' 113 | 114 | # mock_aws_batch.status = 'FAILED' 115 | # pusher.run(dbsession) 116 | # assert pushed_task_instance.status == 'failed' 117 | # assert pushed_task_instance.ended_at == datetime(2017, 6, 4, 6) 118 | 119 | def test_fail_retry(dbsession, monkeypatch): 120 | mock_aws_batch = MockAWSBatch([], 'SUBMITTED') 121 | monkeypatch.setattr(boto3, 'client', mockbatch(mock_aws_batch)) 122 | 123 | task1 = Task(name='task1', active=True, push_destination='aws_batch', retries=1) 124 | dbsession.add(task1) 125 | taskflow = Taskflow() 126 | taskflow.add_task(task1) 127 | taskflow.add_push_worker(AWSBatchPushWorker(taskflow)) 128 | taskflow.sync_db(dbsession) 129 | 130 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 131 | dbsession.add(task_instance) 132 | dbsession.commit() 133 | 134 | pusher = Pusher(taskflow, now_override=datetime(2017, 6, 4, 6)) 135 | pusher.run(dbsession) 136 | 137 | task_instances = dbsession.query(TaskInstance).all() 138 | pushed_task_instance = task_instances[0] 139 | assert pushed_task_instance.status == 'pushed' 140 | assert 'jobId' in pushed_task_instance.push_state 141 | assert pushed_task_instance.push_state['jobName'] == 'task1__1' 142 | 143 | mock_aws_batch.status = 'FAILED' 144 | pusher.run(dbsession) 145 | assert pushed_task_instance.status == 'retry' 146 | assert pushed_task_instance.ended_at == None 147 | 148 | mock_aws_batch.remove_job(pushed_task_instance.push_state['jobId']) 149 | pusher.now_override = datetime(2017, 6, 4, 6, 5) 150 | 151 | mock_aws_batch.status = 'FAILED' 152 | pusher.run(dbsession) 153 | assert pushed_task_instance.status == 'retry' 154 | assert pushed_task_instance.ended_at == None 155 | 156 | mock_aws_batch.remove_job(pushed_task_instance.push_state['jobId']) 157 | pusher.now_override = datetime(2017, 6, 4, 6, 10) 158 | 159 | mock_aws_batch.status = 'FAILED' 160 | pusher.run(dbsession) 161 | assert pushed_task_instance.status == 'failed' 162 | assert pushed_task_instance.ended_at == datetime(2017, 6, 4, 6, 10) 163 | 164 | ## TODO: test task and task_instance job_queue param 165 | 166 | ## TODO: test task and task_instance job_definition param 167 | 168 | ## TODO: assert that all parameters are strings 169 | 170 | ## TODO: assert that all environment values are string 171 | 172 | ## TODO: assert env vars exists, each individually 173 | -------------------------------------------------------------------------------- /tests/test_rest.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from shared_fixtures import * 4 | from restful_ben.test_utils import json_call, login as orig_login, dict_contains, iso_regex 5 | 6 | def login(*args, **kwargs): 7 | kwargs['path'] = '/v1/session' 8 | return orig_login(*args, **kwargs) 9 | 10 | def test_list_workflows(app): 11 | test_client = app.test_client() 12 | login(test_client) 13 | 14 | response = json_call(test_client.get, '/v1/workflows') 15 | assert response.status_code == 200 16 | assert response.json['count'] == 2 17 | assert response.json['page'] == 1 18 | assert response.json['total_pages'] == 1 19 | assert len(response.json['data']) == 2 20 | assert dict_contains(response.json['data'][0], { 21 | 'name': 'workflow1', 22 | 'active': True, 23 | 'title': None, 24 | 'description': None, 25 | 'schedule': '0 6 * * *', 26 | 'start_date': None, 27 | 'end_date': None, 28 | 'concurrency': 1, 29 | 'sla': None, 30 | 'default_priority': 'normal' 31 | }) 32 | assert dict_contains(response.json['data'][1], { 33 | 'name': 'workflow2', 34 | 'active': True, 35 | 'title': None, 36 | 'description': None, 37 | 'schedule': None, 38 | 'start_date': None, 39 | 'end_date': None, 40 | 'concurrency': 1, 41 | 'sla': None, 42 | 'default_priority': 'normal' 43 | }) 44 | 45 | def test_get_workflow(app): 46 | test_client = app.test_client() 47 | login(test_client) 48 | 49 | response = json_call(test_client.get, '/v1/workflows/workflow1') 50 | assert response.status_code == 200 51 | assert dict_contains(response.json, { 52 | 'name': 'workflow1', 53 | 'active': True, 54 | 'title': None, 55 | 'description': None, 56 | 'schedule': '0 6 * * *', 57 | 'start_date': None, 58 | 'end_date': None, 59 | 'concurrency': 1, 60 | 'sla': None, 61 | 'default_priority': 'normal' 62 | }) 63 | 64 | def test_list_tasks(app): 65 | test_client = app.test_client() 66 | login(test_client) 67 | 68 | response = json_call(test_client.get, '/v1/tasks') 69 | assert response.status_code == 200 70 | assert response.json['count'] == 4 71 | assert response.json['page'] == 1 72 | assert response.json['total_pages'] == 1 73 | assert len(response.json['data']) == 4 74 | assert dict_contains(response.json['data'][0], { 75 | 'name': 'task1', 76 | 'workflow_name': 'workflow1', 77 | 'active': True, 78 | 'title': None, 79 | 'description': None, 80 | 'schedule': None, 81 | 'start_date': None, 82 | 'end_date': None, 83 | 'concurrency': 1, 84 | 'sla': None, 85 | 'default_priority': 'normal' 86 | }) 87 | assert dict_contains(response.json['data'][1], { 88 | 'name': 'task2', 89 | 'workflow_name': 'workflow1', 90 | 'active': True, 91 | 'title': None, 92 | 'description': None, 93 | 'schedule': None, 94 | 'start_date': None, 95 | 'end_date': None, 96 | 'concurrency': 1, 97 | 'sla': None, 98 | 'default_priority': 'normal' 99 | }) 100 | assert dict_contains(response.json['data'][2], { 101 | 'name': 'task3', 102 | 'workflow_name': 'workflow1', 103 | 'active': True, 104 | 'title': None, 105 | 'description': None, 106 | 'schedule': None, 107 | 'start_date': None, 108 | 'end_date': None, 109 | 'concurrency': 1, 110 | 'sla': None, 111 | 'default_priority': 'normal' 112 | }) 113 | assert dict_contains(response.json['data'][3], { 114 | 'name': 'task4', 115 | 'workflow_name': 'workflow1', 116 | 'active': True, 117 | 'title': None, 118 | 'description': None, 119 | 'schedule': None, 120 | 'start_date': None, 121 | 'end_date': None, 122 | 'concurrency': 1, 123 | 'sla': None, 124 | 'default_priority': 'normal' 125 | }) 126 | 127 | def test_create_workflow_instance(app, instances): 128 | test_client = app.test_client() 129 | csrf_token = login(test_client) 130 | 131 | workflow_instance = { 132 | 'workflow_name': 'workflow2', 133 | 'unique': 'user-32324-payment-973794' 134 | } 135 | 136 | response = json_call(test_client.post, '/v1/workflow-instances', workflow_instance, headers={'X-CSRF': csrf_token}) 137 | assert response.status_code == 201 138 | assert dict_contains(response.json, { 139 | 'id': 2, 140 | 'workflow_name': 'workflow2', 141 | 'status': 'queued', 142 | 'run_at': iso_regex, 143 | 'unique': 'user-32324-payment-973794', 144 | 'params': None, 145 | 'priority': 'normal', 146 | 'started_at': None, 147 | 'scheduled': False, 148 | 'ended_at': None, 149 | 'created_at': iso_regex, 150 | 'updated_at': iso_regex 151 | }) 152 | 153 | def test_get_workflow_instance(app, instances): 154 | test_client = app.test_client() 155 | login(test_client) 156 | 157 | response = json_call(test_client.get, '/v1/workflow-instances/1') 158 | assert response.status_code == 200 159 | assert dict_contains(response.json, { 160 | 'id': 1, 161 | 'workflow_name': 'workflow1', 162 | 'status': 'running', 163 | 'run_at': '2017-06-03T06:00:00+00:00', 164 | 'unique': None, 165 | 'params': None, 166 | 'priority': 'normal', 167 | 'started_at': '2017-06-03T06:00:00+00:00', 168 | 'scheduled': True, 169 | 'ended_at': None, 170 | 'created_at': iso_regex, 171 | 'updated_at': iso_regex 172 | }) 173 | 174 | def test_list_workflow_instances(app, instances): 175 | test_client = app.test_client() 176 | login(test_client) 177 | 178 | response = json_call(test_client.get, '/v1/workflow-instances') 179 | assert response.status_code == 200 180 | assert response.json['count'] == 1 181 | assert response.json['page'] == 1 182 | assert response.json['total_pages'] == 1 183 | assert len(response.json['data']) == 1 184 | assert dict_contains(response.json['data'][0], { 185 | 'id': 1, 186 | 'workflow_name': 'workflow1', 187 | 'status': 'running', 188 | 'run_at': '2017-06-03T06:00:00+00:00', 189 | 'unique': None, 190 | 'params': None, 191 | 'priority': 'normal', 192 | 'started_at': '2017-06-03T06:00:00+00:00', 193 | 'scheduled': True, 194 | 'ended_at': None, 195 | 'created_at': iso_regex, 196 | 'updated_at': iso_regex 197 | }) 198 | 199 | def test_update_workflow_instance(app, instances): 200 | test_client = app.test_client() 201 | csrf_token = login(test_client) 202 | 203 | response = json_call(test_client.get, '/v1/workflow-instances/1') 204 | assert response.status_code == 200 205 | 206 | workflow_instance = response.json 207 | workflow_instance['priority'] = 'high' 208 | previous_updated_at = response.json['updated_at'] 209 | 210 | response = json_call(test_client.put, '/v1/workflow-instances/1', workflow_instance, headers={'X-CSRF': csrf_token}) 211 | assert response.status_code == 200 212 | assert dict_contains(response.json, { 213 | 'id': 1, 214 | 'workflow_name': 'workflow1', 215 | 'status': 'running', 216 | 'run_at': iso_regex, 217 | 'unique': None, 218 | 'params': None, 219 | 'priority': 'high', 220 | 'started_at': iso_regex, 221 | 'scheduled': True, 222 | 'ended_at': None, 223 | 'created_at': iso_regex, 224 | 'updated_at': iso_regex 225 | }) 226 | assert response.json['updated_at'] > previous_updated_at 227 | 228 | def test_delete_workflow_instance(app, instances): 229 | test_client = app.test_client() 230 | csrf_token = login(test_client) 231 | 232 | response = json_call(test_client.get, '/v1/workflow-instances/1') 233 | assert response.status_code == 200 234 | 235 | response = json_call(test_client.get, '/v1/task-instances?workflow_instance_id=1') 236 | assert response.status_code == 200 237 | assert response.json['count'] == 4 238 | 239 | response = json_call(test_client.delete, '/v1/workflow-instances/1', headers={'X-CSRF': csrf_token}) 240 | assert response.status_code == 204 241 | 242 | response = json_call(test_client.get, '/v1/workflow-instances/1') 243 | assert response.status_code == 404 244 | 245 | response = json_call(test_client.get, '/v1/task-instances?workflow_instance_id=1') 246 | assert response.status_code == 200 247 | assert response.json['count'] == 0 248 | 249 | def test_create_task_instance(app, instances): 250 | test_client = app.test_client() 251 | csrf_token = login(test_client) 252 | 253 | task_instance = { 254 | 'task_name': 'task1', 255 | 'unique': 'user-32324-payment-973794' 256 | } 257 | 258 | response = json_call(test_client.post, '/v1/task-instances', task_instance, headers={'X-CSRF': csrf_token}) 259 | assert response.status_code == 201 260 | assert dict_contains(response.json, { 261 | 'id': 5, 262 | 'task_name': 'task1', 263 | 'workflow_instance_id': None, 264 | 'status': 'queued', 265 | 'run_at': iso_regex, 266 | 'unique': 'user-32324-payment-973794', 267 | 'params': {}, 268 | 'priority': 'normal', 269 | 'started_at': None, 270 | 'scheduled': False, 271 | 'ended_at': None, 272 | 'attempts': 0, 273 | 'max_attempts': 1, 274 | 'timeout': 300, 275 | 'retry_delay': 300, 276 | 'push': False, 277 | 'push_state': None, 278 | 'worker_id': None, 279 | 'locked_at': None, 280 | 'created_at': iso_regex, 281 | 'updated_at': iso_regex 282 | }) 283 | 284 | def test_get_task_instance(app, instances): 285 | test_client = app.test_client() 286 | login(test_client) 287 | 288 | response = json_call(test_client.get, '/v1/task-instances/1') 289 | assert response.status_code == 200 290 | assert dict_contains(response.json, { 291 | 'id': 1, 292 | 'task_name': 'task1', 293 | 'workflow_instance_id': 1, 294 | 'status': 'success', 295 | 'run_at': iso_regex, 296 | 'unique': None, 297 | 'params': {}, 298 | 'priority': 'normal', 299 | 'started_at': iso_regex, 300 | 'scheduled': True, 301 | 'ended_at': iso_regex, 302 | 'attempts': 1, 303 | 'max_attempts': 1, 304 | 'timeout': 300, 305 | 'retry_delay': 300, 306 | 'push': False, 307 | 'push_state': None, 308 | 'worker_id': None, 309 | 'locked_at': None, 310 | 'created_at': iso_regex, 311 | 'updated_at': iso_regex 312 | }) 313 | 314 | def test_list_task_instances(app, instances): 315 | test_client = app.test_client() 316 | login(test_client) 317 | 318 | response = json_call(test_client.get, '/v1/task-instances') 319 | assert response.status_code == 200 320 | assert response.json['count'] == 4 321 | assert response.json['page'] == 1 322 | assert response.json['total_pages'] == 1 323 | assert len(response.json['data']) == 4 324 | assert dict_contains(response.json['data'][0], { 325 | 'id': 1, 326 | 'task_name': 'task1', 327 | 'workflow_instance_id': 1, 328 | 'status': 'success', 329 | 'run_at': iso_regex, 330 | 'unique': None, 331 | 'params': {}, 332 | 'priority': 'normal', 333 | 'started_at': iso_regex, 334 | 'scheduled': True, 335 | 'ended_at': iso_regex, 336 | 'attempts': 1, 337 | 'max_attempts': 1, 338 | 'timeout': 300, 339 | 'retry_delay': 300, 340 | 'push': False, 341 | 'push_state': None, 342 | 'worker_id': None, 343 | 'locked_at': None, 344 | 'created_at': iso_regex, 345 | 'updated_at': iso_regex 346 | }) 347 | assert dict_contains(response.json['data'][1], { 348 | 'id': 2, 349 | 'task_name': 'task2', 350 | 'workflow_instance_id': 1, 351 | 'status': 'success', 352 | 'run_at': iso_regex, 353 | 'unique': None, 354 | 'params': {}, 355 | 'priority': 'normal', 356 | 'started_at': iso_regex, 357 | 'scheduled': True, 358 | 'ended_at': iso_regex, 359 | 'attempts': 1, 360 | 'max_attempts': 1, 361 | 'timeout': 300, 362 | 'retry_delay': 300, 363 | 'push': False, 364 | 'push_state': None, 365 | 'worker_id': None, 366 | 'locked_at': None, 367 | 'created_at': iso_regex, 368 | 'updated_at': iso_regex 369 | }) 370 | assert dict_contains(response.json['data'][2], { 371 | 'id': 3, 372 | 'task_name': 'task3', 373 | 'workflow_instance_id': 1, 374 | 'status': 'success', 375 | 'run_at': iso_regex, 376 | 'unique': None, 377 | 'params': {}, 378 | 'priority': 'normal', 379 | 'started_at': iso_regex, 380 | 'scheduled': True, 381 | 'ended_at': iso_regex, 382 | 'attempts': 1, 383 | 'max_attempts': 1, 384 | 'timeout': 300, 385 | 'retry_delay': 300, 386 | 'push': False, 387 | 'push_state': None, 388 | 'worker_id': None, 389 | 'locked_at': None, 390 | 'created_at': iso_regex, 391 | 'updated_at': iso_regex 392 | }) 393 | assert dict_contains(response.json['data'][3], { 394 | 'id': 4, 395 | 'task_name': 'task4', 396 | 'workflow_instance_id': 1, 397 | 'status': 'running', 398 | 'run_at': iso_regex, 399 | 'unique': None, 400 | 'params': {}, 401 | 'priority': 'normal', 402 | 'started_at': iso_regex, 403 | 'scheduled': True, 404 | 'ended_at': None, 405 | 'attempts': 1, 406 | 'max_attempts': 1, 407 | 'timeout': 300, 408 | 'retry_delay': 300, 409 | 'push': False, 410 | 'push_state': None, 411 | 'worker_id': None, 412 | 'locked_at': None, 413 | 'created_at': iso_regex, 414 | 'updated_at': iso_regex 415 | }) 416 | 417 | def test_update_task_instance(app, instances): 418 | test_client = app.test_client() 419 | csrf_token = login(test_client) 420 | 421 | response = json_call(test_client.get, '/v1/task-instances/1') 422 | assert response.status_code == 200 423 | 424 | task_instance = response.json 425 | task_instance['worker_id'] = 'foo' 426 | previous_updated_at = task_instance['updated_at'] 427 | 428 | response = json_call(test_client.put, '/v1/task-instances/1', task_instance, headers={'X-CSRF': csrf_token}) 429 | assert response.status_code == 200 430 | assert dict_contains(response.json, { 431 | 'id': 1, 432 | 'task_name': 'task1', 433 | 'workflow_instance_id': 1, 434 | 'status': 'success', 435 | 'run_at': iso_regex, 436 | 'unique': None, 437 | 'params': {}, 438 | 'priority': 'normal', 439 | 'started_at': iso_regex, 440 | 'scheduled': True, 441 | 'ended_at': iso_regex, 442 | 'attempts': 1, 443 | 'max_attempts': 1, 444 | 'timeout': 300, 445 | 'retry_delay': 300, 446 | 'push': False, 447 | 'push_state': None, 448 | 'worker_id': 'foo', 449 | 'locked_at': None, 450 | 'created_at': iso_regex, 451 | 'updated_at': iso_regex 452 | }) 453 | assert response.json['updated_at'] > previous_updated_at 454 | 455 | def test_delete_task_instance(app, instances): 456 | test_client = app.test_client() 457 | csrf_token = login(test_client) 458 | 459 | response = json_call(test_client.get, '/v1/task-instances/1') 460 | assert response.status_code == 200 461 | 462 | response = json_call(test_client.delete, '/v1/task-instances/1', headers={'X-CSRF': csrf_token}) 463 | assert response.status_code == 204 464 | 465 | response = json_call(test_client.get, '/v1/task-instances/1') 466 | assert response.status_code == 404 467 | 468 | def test_get_recurring_latest(app, instances, dbsession): 469 | test_client = app.test_client() 470 | login(test_client) 471 | 472 | workflow_instance = WorkflowInstance( 473 | workflow_name='workflow1', 474 | scheduled=True, 475 | run_at=datetime(2017, 6, 4, 6), 476 | started_at=datetime(2017, 6, 4, 6), 477 | status='running', 478 | priority='normal') 479 | dbsession.add(workflow_instance) 480 | dbsession.commit() 481 | 482 | response = json_call(test_client.get, '/v1/workflow-instances/recurring-latest') 483 | assert response.status_code == 200 484 | assert dict_contains(response.json[0], { 485 | 'id': 2, 486 | 'workflow_name': 'workflow1', 487 | 'status': 'running', 488 | 'run_at': '2017-06-04T06:00:00+00:00', 489 | 'unique': None, 490 | 'params': None, 491 | 'priority': 'normal', 492 | 'started_at': '2017-06-04T06:00:00+00:00', 493 | 'scheduled': True, 494 | 'ended_at': None, 495 | 'created_at': iso_regex, 496 | 'updated_at': iso_regex 497 | }) 498 | -------------------------------------------------------------------------------- /tests/test_tasks.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | 3 | import pytest 4 | from sqlalchemy.exc import IntegrityError 5 | 6 | from taskflow import Scheduler, Taskflow, Task, TaskInstance 7 | from shared_fixtures import * 8 | 9 | get_logging() 10 | 11 | @pytest.fixture 12 | def tasks(dbsession): 13 | task1 = Task(name='task1', active=True) 14 | task2 = Task(name='task2', active=True, schedule='0 6 * * *') 15 | task3 = Task(name='task3', active=True) 16 | task4 = Task(name='task4', active=True, schedule='0 2 * * *') 17 | 18 | dbsession.add(task1) 19 | dbsession.add(task2) 20 | dbsession.add(task3) 21 | dbsession.add(task4) 22 | dbsession.commit() 23 | 24 | return [task1, task2, task3, task4] 25 | 26 | ## TODO: test multiple works pulling, one wins 27 | 28 | ## TODO: test task inactive 29 | 30 | ## TODO: test touching / relocking? 31 | 32 | ## TODO: test passing task_names 33 | 34 | def test_schedule_recurring_task(dbsession, tasks): 35 | taskflow = Taskflow() 36 | taskflow.add_tasks(tasks) 37 | scheduler = Scheduler(taskflow, now_override=datetime(2017, 6, 3, 6)) 38 | scheduler.run(dbsession) 39 | 40 | task_instances = dbsession.query(TaskInstance).all() 41 | 42 | ## TODO: how was this commited before session.commit() was added queue_task ? 43 | ## TODO: ^ test a future run_at 44 | 45 | assert len(task_instances) == 2 46 | for task_instance in task_instances: 47 | assert task_instance.status == 'queued' 48 | assert task_instance.scheduled == True 49 | 50 | if task_instance.task_name == 'task2': 51 | assert task_instance.run_at == datetime(2017, 6, 4, 6) 52 | else: 53 | assert task_instance.run_at == datetime(2017, 6, 4, 2) 54 | 55 | def test_queue_unique(dbsession, engine): 56 | task1 = Task(name='task1', active=True) 57 | dbsession.add(task1) 58 | taskflow = Taskflow() 59 | taskflow.add_task(task1) 60 | 61 | task_instance1 = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6), unique='foo') 62 | dbsession.add(task_instance1) 63 | dbsession.commit() 64 | 65 | with pytest.raises(IntegrityError): 66 | task_instance2 = task1.get_new_instance(run_at=datetime(2017, 6, 4, 8), unique='foo') 67 | dbsession.add(task_instance2) 68 | dbsession.commit() 69 | 70 | def test_queue_pull_task(dbsession, engine): 71 | task1 = Task(name='task1', active=True) 72 | dbsession.add(task1) 73 | taskflow = Taskflow() 74 | taskflow.add_task(task1) 75 | 76 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 77 | dbsession.add(task_instance) 78 | dbsession.commit() 79 | 80 | task_instance_id = task_instance.id 81 | dbsession.expunge_all() 82 | 83 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 6, 0, 12)) 84 | pulled_task_instance = pulled_task_instances[0] 85 | 86 | assert pulled_task_instance.id == task_instance_id 87 | assert pulled_task_instance.status == 'running' 88 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 89 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 90 | assert pulled_task_instance.worker_id == 'test' 91 | 92 | def test_queue_pull_task_priority(dbsession, engine): 93 | task1 = Task(name='task1', active=True) 94 | dbsession.add(task1) 95 | taskflow = Taskflow() 96 | taskflow.add_task(task1) 97 | 98 | task_instance1 = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 99 | task_instance2 = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6), priority='high') 100 | dbsession.add(task_instance1) 101 | dbsession.add(task_instance2) 102 | dbsession.commit() 103 | 104 | normal_task_instance_id = task_instance1.id 105 | high_task_instance_id = task_instance2.id 106 | 107 | dbsession.expunge_all() 108 | 109 | pulled_task_instances = taskflow.pull(dbsession, 'test', max_tasks=1, now=datetime(2017, 6, 4, 6, 0, 12)) 110 | pulled_task_instance = pulled_task_instances[0] 111 | 112 | assert pulled_task_instance.id == high_task_instance_id 113 | assert pulled_task_instance.status == 'running' 114 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 115 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 116 | assert pulled_task_instance.worker_id == 'test' 117 | assert pulled_task_instance.priority == 'high' 118 | 119 | pulled_task_instances = taskflow.pull(dbsession, 'test', max_tasks=1, now=datetime(2017, 6, 4, 6, 0, 12)) 120 | pulled_task_instance = pulled_task_instances[0] 121 | 122 | assert pulled_task_instance.id == normal_task_instance_id 123 | assert pulled_task_instance.status == 'running' 124 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 125 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 126 | assert pulled_task_instance.worker_id == 'test' 127 | assert pulled_task_instance.priority == 'normal' 128 | 129 | def test_queue_pull_task_run_at_order(dbsession, engine): 130 | task1 = Task(name='task1', active=True) 131 | dbsession.add(task1) 132 | taskflow = Taskflow() 133 | taskflow.add_task(task1) 134 | 135 | task_instance1 = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6, 0, 5)) 136 | task_instance2 = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6, 0, 10)) 137 | dbsession.add(task_instance1) 138 | dbsession.add(task_instance2) 139 | dbsession.commit() 140 | 141 | sooner_task_instance_id = task_instance1.id 142 | later_task_instance_id = task_instance2.id 143 | 144 | dbsession.expunge_all() 145 | 146 | pulled_task_instances = taskflow.pull(dbsession, 'test', max_tasks=1, now=datetime(2017, 6, 4, 6, 0, 12)) 147 | pulled_task_instance = pulled_task_instances[0] 148 | 149 | assert pulled_task_instance.id == sooner_task_instance_id 150 | assert pulled_task_instance.status == 'running' 151 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 152 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 153 | assert pulled_task_instance.worker_id == 'test' 154 | 155 | pulled_task_instances = taskflow.pull(dbsession, 'test', max_tasks=1, now=datetime(2017, 6, 4, 6, 0, 12)) 156 | pulled_task_instance = pulled_task_instances[0] 157 | 158 | assert pulled_task_instance.id == later_task_instance_id 159 | assert pulled_task_instance.status == 'running' 160 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 161 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 162 | assert pulled_task_instance.worker_id == 'test' 163 | 164 | def test_succeed_task(dbsession, engine): 165 | task1 = Task(name='task1', active=True) 166 | dbsession.add(task1) 167 | taskflow = Taskflow() 168 | taskflow.add_task(task1) 169 | 170 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 171 | dbsession.add(task_instance) 172 | dbsession.commit() 173 | 174 | task_instance_id = task_instance.id 175 | dbsession.expunge_all() 176 | 177 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 6, 0, 12)) 178 | pulled_task_instance = pulled_task_instances[0] 179 | 180 | assert pulled_task_instance.id == task_instance_id 181 | assert pulled_task_instance.status == 'running' 182 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 183 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 184 | assert pulled_task_instance.worker_id == 'test' 185 | 186 | pulled_task_instance.succeed(dbsession, taskflow, now=datetime(2017, 6, 4, 6, 0, 15)) 187 | dbsession.refresh(pulled_task_instance) 188 | 189 | assert pulled_task_instance.status == 'success' 190 | assert pulled_task_instance.ended_at == datetime(2017, 6, 4, 6, 0, 15) 191 | 192 | def test_fail_task(dbsession, engine): 193 | task1 = Task(name='task1', active=True) 194 | dbsession.add(task1) 195 | taskflow = Taskflow() 196 | taskflow.add_task(task1) 197 | 198 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 199 | dbsession.add(task_instance) 200 | dbsession.commit() 201 | 202 | task_instance_id = task_instance.id 203 | dbsession.expunge_all() 204 | 205 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 6, 0, 12)) 206 | pulled_task_instance = pulled_task_instances[0] 207 | 208 | assert pulled_task_instance.id == task_instance_id 209 | assert pulled_task_instance.status == 'running' 210 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 211 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 212 | assert pulled_task_instance.worker_id == 'test' 213 | 214 | pulled_task_instance.fail(dbsession, taskflow, now=datetime(2017, 6, 4, 6, 0, 15)) 215 | dbsession.refresh(pulled_task_instance) 216 | 217 | assert pulled_task_instance.status == 'failed' 218 | assert pulled_task_instance.ended_at == datetime(2017, 6, 4, 6, 0, 15) 219 | 220 | def test_timeout_task(dbsession, engine): 221 | task1 = Task(name='task1', active=True, retries=2) 222 | dbsession.add(task1) 223 | taskflow = Taskflow() 224 | taskflow.add_task(task1) 225 | 226 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 227 | dbsession.add(task_instance) 228 | dbsession.commit() 229 | 230 | task_instance_id = task_instance.id 231 | dbsession.expunge_all() 232 | 233 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 6, 0, 12)) 234 | pulled_task_instance = pulled_task_instances[0] 235 | 236 | assert pulled_task_instance.id == task_instance_id 237 | assert pulled_task_instance.status == 'running' 238 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 239 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 240 | assert pulled_task_instance.worker_id == 'test' 241 | 242 | dbsession.expunge_all() 243 | 244 | ## pull 5 minutes later 245 | pulled_task_instances = taskflow.pull(dbsession, 'test2', now=datetime(2017, 6, 4, 6, 5, 15)) 246 | pulled_task_instance = pulled_task_instances[0] 247 | assert pulled_task_instance.status == 'running' 248 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 5, 15) 249 | assert pulled_task_instance.worker_id == 'test2' 250 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 251 | assert pulled_task_instance.attempts == 2 252 | 253 | def test_custom_timeout_task(dbsession, engine): 254 | task1 = Task(name='task1', active=True, retries=2, timeout=7200) 255 | dbsession.add(task1) 256 | taskflow = Taskflow() 257 | taskflow.add_task(task1) 258 | 259 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 260 | dbsession.add(task_instance) 261 | dbsession.commit() 262 | 263 | task_instance_id = task_instance.id 264 | dbsession.expunge_all() 265 | 266 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 6, 0, 12)) 267 | pulled_task_instance = pulled_task_instances[0] 268 | 269 | assert pulled_task_instance.id == task_instance_id 270 | assert pulled_task_instance.status == 'running' 271 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 0, 12) 272 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 273 | assert pulled_task_instance.worker_id == 'test' 274 | 275 | dbsession.expunge_all() 276 | 277 | ## pull almost 2 hours later 278 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 8, 0, 10)) 279 | assert pulled_task_instances == [] 280 | 281 | ## pull after 2 hours later 282 | pulled_task_instances = taskflow.pull(dbsession, 'test2', now=datetime(2017, 6, 4, 8, 0, 15)) 283 | pulled_task_instance = pulled_task_instances[0] 284 | assert pulled_task_instance.status == 'running' 285 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 8, 0, 15) 286 | assert pulled_task_instance.worker_id == 'test2' 287 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 288 | assert pulled_task_instance.attempts == 2 289 | 290 | def test_task_retry_success(dbsession, engine): 291 | task1 = Task(name='task1', active=True, retries=1) 292 | dbsession.add(task1) 293 | taskflow = Taskflow() 294 | taskflow.add_task(task1) 295 | 296 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 297 | dbsession.add(task_instance) 298 | dbsession.commit() 299 | 300 | task_instance_id = task_instance.id 301 | dbsession.expunge_all() 302 | 303 | now = datetime(2017, 6, 4, 6, 0, 12) 304 | 305 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=now) 306 | pulled_task_instance = pulled_task_instances[0] 307 | 308 | assert pulled_task_instance.id == task_instance_id 309 | assert pulled_task_instance.status == 'running' 310 | assert pulled_task_instance.locked_at == now 311 | assert pulled_task_instance.started_at == now 312 | assert pulled_task_instance.worker_id == 'test' 313 | 314 | pulled_task_instance.fail(dbsession, taskflow, now=datetime(2017, 6, 4, 6, 0, 15)) 315 | dbsession.refresh(pulled_task_instance) 316 | 317 | assert pulled_task_instance.status == 'retry' 318 | 319 | dbsession.expunge_all() 320 | 321 | ## stays within retry_delay 322 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=datetime(2017, 6, 4, 6, 5, 14)) 323 | assert pulled_task_instances == [] 324 | 325 | pulled_task_instances = taskflow.pull(dbsession, 'test2', now=datetime(2017, 6, 4, 6, 5, 16)) 326 | pulled_task_instance = pulled_task_instances[0] 327 | assert pulled_task_instance.status == 'running' 328 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 5, 16) 329 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 330 | assert pulled_task_instance.attempts == 2 331 | assert pulled_task_instance.worker_id == 'test2' 332 | 333 | pulled_task_instance.succeed(dbsession, taskflow, now=datetime(2017, 6, 4, 6, 5, 20)) 334 | dbsession.refresh(pulled_task_instance) 335 | 336 | assert pulled_task_instance.status == 'success' 337 | assert pulled_task_instance.ended_at == datetime(2017, 6, 4, 6, 5, 20) 338 | 339 | def test_task_retry_fail(dbsession, engine): 340 | task1 = Task(name='task1', active=True, retries=1) 341 | dbsession.add(task1) 342 | taskflow = Taskflow() 343 | taskflow.add_task(task1) 344 | 345 | task_instance = task1.get_new_instance(run_at=datetime(2017, 6, 4, 6)) 346 | dbsession.add(task_instance) 347 | dbsession.commit() 348 | 349 | task_instance_id = task_instance.id 350 | dbsession.expunge_all() 351 | 352 | now = datetime(2017, 6, 4, 6, 0, 12) 353 | 354 | pulled_task_instances = taskflow.pull(dbsession, 'test', now=now) 355 | pulled_task_instance = pulled_task_instances[0] 356 | 357 | assert pulled_task_instance.id == task_instance_id 358 | assert pulled_task_instance.status == 'running' 359 | assert pulled_task_instance.locked_at == now 360 | assert pulled_task_instance.started_at == now 361 | assert pulled_task_instance.worker_id == 'test' 362 | 363 | pulled_task_instance.fail(dbsession, taskflow, now=datetime(2017, 6, 4, 6, 0, 15)) 364 | dbsession.refresh(pulled_task_instance) 365 | 366 | assert pulled_task_instance.status == 'retry' 367 | 368 | dbsession.expunge_all() 369 | 370 | pulled_task_instances = taskflow.pull(dbsession, 'test2', now=datetime(2017, 6, 4, 6, 5, 16)) 371 | pulled_task_instance = pulled_task_instances[0] 372 | assert pulled_task_instance.status == 'running' 373 | assert pulled_task_instance.locked_at == datetime(2017, 6, 4, 6, 5, 16) 374 | assert pulled_task_instance.started_at == datetime(2017, 6, 4, 6, 0, 12) 375 | assert pulled_task_instance.attempts == 2 376 | assert pulled_task_instance.worker_id == 'test2' 377 | 378 | pulled_task_instance.fail(dbsession, taskflow, now=datetime(2017, 6, 4, 6, 5, 20)) 379 | dbsession.refresh(pulled_task_instance) 380 | 381 | assert pulled_task_instance.status == 'failed' 382 | assert pulled_task_instance.ended_at == datetime(2017, 6, 4, 6, 5, 20) 383 | --------------------------------------------------------------------------------