Importing taskflow functions from a library

[MatrixManAtYrService]

I tried to make a library which allows users to consume @task decorated functions. It didn't go so well. Here's an example of why:

# lib.py
@task
def incr(x):
    return x + 1

# dag.py
from lib import incr

@dag(...)
def my_dag():
    incr(incr(1))  

This works fine, until you want to set something like retries=2 on the task. To do this you have to modify the decorator, which is in lib.py. That's a problem because I don't want to modify it for all consumers, just for this instance.

I see a few possible approaches to working around this:

• don't use the taskflow interface, put custom operators in the lib instead. User can't use function composition to express dependency 🙁

• instead of importing decorated functions, import functions that return decorated functions. The caller might do something like this:
  incr_factory(retries=2)(incr_factory(retries=2)(1))

• modify airlflow to detect that retries is not a kwarg on incr. In this case I'm using the airflow_ prefix for this:
  incr(incr(1, airflow_retries=2), airflow_retries=2)

If there are best practices here I want to be sure I'm using them in my library, so my question is: Is there a right way to do this? If not, how would you do it?

[josh-fell]

Hey @MatrixManAtYrService! With 2.3, there was a new override() method added to decorators which might help here. You are now able to assign task-level (BaseOperator-level) params without directly setting them in the decorator explicitly.

Your example could now be written as:

# lib.py
from airflow.decorators import task

@task
def incr(x):
    return x + 1

# dag.py
from include.lib import incr
from pendulum import datetime

from airflow.decorators import dag

@dag(start_date=datetime(2022, 1, 1), schedule_interval=None)
def my_dag():
    first_incr = incr.override(retries=5)(1)
    incr.override(retries=2)(first_incr)

_ = my_dag()

Of course, if you didn't want your lib functions to require task-decoration (to keep your business logic and Airflow-domain separate for example), you can do:

# lib.py
def incr(x):
    return x + 1

# dag.py
from include.lib import incr
from pendulum import datetime

from airflow.decorators import dag, task

@dag(start_date=datetime(2022, 1, 1), schedule_interval=None)
def my_dag():
    first_incr = task(incr, retries=5)(1)
    task(incr, retries=2)(first_incr)

_ = my_dag()

[MatrixManAtYrService]

Both of those solutions look like they scratch the itch, thank you!