Airflow cluster policy

airflow_local_settings.py

"""Local Airflow settings (policy)"""

from datetime import timedelta

from airflow.models import TaskInstance

DEFAULT_EXECUTION_TIMEOUT = timedelta(hours=48)


def policy(task: TaskInstance) -> None:
    """Alter tasks before they get run.

    :param task: task instance that will be executed
    """
    print("cluster policy got task", task)  # debug
    if hasattr(task, "execution_timeout") and not task.execution_timeout:
        print("Setting execution timeout on task '{} ({})' to '{}'.".format(
            task.task_id,
            task.task_type,
            DEFAULT_EXECUTION_TIMEOUT
        ))
        task.execution_timeout = DEFAULT_EXECUTION_TIMEOUT