<filename>bd-runtime/airflow/dags/test_short_circ.py
from airflow import DAG
from airflow.models import SkipMixin, BaseOperator, TaskInstance
from airflow.operators.bash import BashOperator
from airflow.operators.dummy import DummyOperator
from airflow.operators.python import PythonOperator
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.providers.apache.livy.operators.livy import LivyOperator
from airflow.utils import timezone
from airflow.utils.dates import days_ago
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.state import State
from airflow.exceptions import AirflowFailException
from typing import Dict
from airflow.utils.session import create_session


class FinishOperator(BaseOperator, SkipMixin):

    def execute(self, context: Dict):
        excl_states = [State.SUCCESS, State.FAILED, State.SKIPPED, State.RUNNING]
        all_tasks = [t for t in context['dag_run'].get_task_instances() if t.state not in excl_states]
        if all_tasks:
            self.log.info('Skipping tasks: [%s]',
                          ', '.join([t.task_id + '(' + str(t.state) + ')' for t in all_tasks]))
            self.skip(context['dag_run'], context['ti'].execution_date, all_tasks)


class KillOperator(FinishOperator):

    def __init__(self, message, task_id: str, **kwargs) -> None:
        super().__init__(task_id=task_id, **kwargs)
        self.message = message

    def execute(self, context: Dict):
        super().execute(context)
        # self.log.info(self.message)
        raise AirflowFailException(self.message)


def on_failure_callback(context):
    ti = context['task_instance']
    task = ti.task
    downstream_tasks = task.downstream_list
    branch_rules = {TriggerRule.ONE_FAILED, TriggerRule.ALL_FAILED}
    branch_task_ids = {t.task_id for t in downstream_tasks if t.trigger_rule in branch_rules}
    for branch_task_id in list(branch_task_ids):
        branch_task_ids.update(task.dag.get_task(branch_task_id).get_flat_relative_ids(upstream=False))
    skip_tasks_ids = [t.task_id for t in downstream_tasks if t.task_id not in branch_task_ids]
    task.log.info("Skipping branches %s", skip_tasks_ids)
    if skip_tasks_ids:
        dag_run = ti.get_dagrun()
        now = timezone.utcnow()
        with create_session() as session:
            session.query(TaskInstance).filter(
                TaskInstance.dag_id == dag_run.dag_id,
                TaskInstance.execution_date == dag_run.execution_date,
                TaskInstance.task_id.in_(skip_tasks_ids),
                ).update(
                {
                    TaskInstance.state: State.SKIPPED,
                    TaskInstance.start_date: now,
                    TaskInstance.end_date: now,
                },
                synchronize_session=False,
            )
            session.commit()


def raise_error(msg):
    raise Exception(msg)


default_args = {
    # 'owner': 'airflow',
    # 'depends_on_past': False,
    # 'email': ['<EMAIL>'],
    # 'email_on_failure': False,
    # 'email_on_retry': False,
    # 'retries': 1,
    # 'retry_delay': timedelta(minutes=5),
    # 'queue': 'bash_queue',
    # 'pool': 'backfill',
    # 'priority_weight': 10,
    # 'end_date': datetime(2016, 1, 1),
    # 'wait_for_downstream': False,
    # 'dag': dag,
    # 'sla': timedelta(hours=2),
    # 'execution_timeout': timedelta(seconds=300),
    # 'on_failure_callback': some_function,
    # 'on_success_callback': some_other_function,
    # 'on_retry_callback': another_function,
    # 'sla_miss_callback': yet_another_function,
    # 'trigger_rule': 'all_success'
}
with DAG(
        'test_short_circ',
        default_args=default_args,
        description='A simple test error handling',
        schedule_interval=None,
        start_date=days_ago(0),
        tags=['test'],
) as dag:
    start = DummyOperator(
        task_id='start'
    )
    a = PythonOperator(
        task_id='a',
        python_callable=lambda: print('a'),
        on_failure_callback=on_failure_callback
    )
    start >> a
    err1 = KillOperator(
        task_id='err1',
        message='a failed',
        trigger_rule=TriggerRule.ALL_FAILED
    )
    a >> err1
    fork = DummyOperator(
        task_id='fork'
    )
    a >> fork
    b = PythonOperator(
        task_id='b',
        python_callable=lambda: print('b'),
        on_failure_callback=on_failure_callback
    )
    fork >> b
    err2 = KillOperator(
        task_id='err2',
        message='b failed',
        trigger_rule=TriggerRule.ONE_FAILED
    )
    b >> err2
    c = PythonOperator(
        task_id='c',
        python_callable=lambda: print('c'),
        on_failure_callback=on_failure_callback
    )
    fork >> c
    err3 = KillOperator(
        task_id='err3',
        message='c failed',
        trigger_rule=TriggerRule.ONE_FAILED
    )
    c >> err3
    join = DummyOperator(
        task_id='join'
    )
    b >> join
    c >> join
    transformation = LivyOperator(
        task_id='aws_test2',
        dag=dag,
        livy_conn_id='livy_default',
        file='s3a://datagram/user/root/deployments/autogenerated_tr_aws_test_2/'
             'ru.neoflex.meta.etl2.spark.aws_test_2-1.0-SNAPSHOT.jar',
        proxy_user='root',
        args=[
            'HOME=/user',
            'USER=root',
            'WF_HOME=/user/root',
            'ROOT_WORKFLOW_ID=aws_test_2_729508210',
            'CURRENT_WORKFLOW_ID=aws_test_2_729508210',
            'SLIDE_SIZE=400',
            'FETCH_SIZE=1000',
            'PARTITION_NUM=1',
            'FAIL_THRESHOLD=1000',
            'DEBUG=true',
            'MASTER='
        ],
        num_executors=1,
        conf={
            'spark.shuffle.compress': 'false',
        },
        class_name='ru.neoflex.meta.etl2.spark.aws_test_2Job',
        polling_interval=5,
        on_failure_callback=on_failure_callback
    )
    join >> transformation
    err4 = KillOperator(
        task_id='err4',
        message='transformation or workflow failed',
        trigger_rule=TriggerRule.ONE_FAILED
    )
    transformation >> err4
    workflow = TriggerDagRunOperator(
        task_id='wf_aws_test',
        trigger_dag_id='wf_aws_test',
        wait_for_completion=True,
        poke_interval=5,
        dag=dag,
        on_failure_callback=on_failure_callback
    )
    transformation >> workflow
    workflow >> err4
    finish = FinishOperator(
        task_id='finish',
    )
    workflow >> finish
