Assign airflow task to several DAGs

Viewed 765

I am trying to reuse an existing airflow task by assigning it to different dags.

def create_new_task_for_dag(task: BaseOperator,
                        dag: models.DAG) -> BaseOperator:
    """Create a deep copy of given task and associate it with given dag
    """
    new_task = copy.deepcopy(task)
    new_task.dag = dag
    return new_task


print_datetime_task = python_operator.PythonOperator(
    task_id='print_datetime', python_callable=_print_datetime)

# define a new dag ...
# add to the new dag
create_new_task_for_dag(print_datetime_task, new_dag)

Then it gives the error Task is missing the start_date parameter. If I define the dag when creating the operator, print_datetime_task = PythonOperator(task_id='print_datetime', python_callable=_print_datetime, dag=new_dag), then it is OK.

I have searched around, and this seems to be the root cause: https://github.com/apache/airflow/pull/5598, but PR has been marked as stale.

I wonder if there is any other approach to reuse an existing airflow task assign to a different dag.

I am using apache-airflow[docker,kubernetes]==1.10.10

1 Answers

While I don't know the solution to your problem with current design (code-layout), it can be made to work by tweaking the design slightly (note that the following code-snippets have NOT been tested)


Instead of copying a task from a DAG,

 def create_new_task_for_dag(task: BaseOperator,
                        dag: models.DAG) -> BaseOperator:
    """Create a deep copy of given task and associate it with given dag
    """
    new_task = copy.deepcopy(task)
    new_task.dag = dag
    return new_task

you can move the instantiation of task (as well as it's assignment to the DAG) to a separate utility function.

from datetime import datetime
from typing import Dict, Any

from airflow.models.dag import DAG
from airflow.operators.python_operator import PythonOperator


def add_new_print_datetime_task(my_dag: DAG,
                                kwargs: Dict[str, Any]) -> PythonOperator:
    """
    Creates and adds a new 'print_datetime' (PythonOperator) task in 'my_dag'
    and returns it's reference
    :param my_dag: reference to DAG object in which to add the task
    :type my_dag: DAG
    :param kwargs: dictionary of args for PythonOperator / BaseOperator
                   'task_id' is mandatory
    :type kwargs: Dict[str, Any]
    :return: PythonOperator
    """

    def my_callable() -> None:
        print(datetime.now())

    return PythonOperator(dag=my_dag, python_callable=my_callable, **kwargs)

Thereafter you can call that function everytime you want to instantiate that same task (and assign to any DAG)

with DAG(dag_id="my_dag_id", start_date=datetime(year=2020, month=8, day=22, hour=16, minute=30)) as my_dag:
    print_datetime_task_kwargs: Dict[str, Any] = {
        "task_id": "my_task_id",
        "depends_on_past": True
    }
    print_datetime_task: PythonOperator = add_new_print_datetime_task(my_dag=my_dag, kwargs=print_datetime_task_kwargs)
    # ... other tasks and their wiring

References / good reads

Related