[AIP-31] Add lazy way to get Airflow context from a wrapped function

[casassg]

Description

Currently the only way to access the context object from a PythonOperator wrapped function is by setting provide_context=True in the Operator. That loads it into the kwargs. This makes the signature of the function become difficult to argue about.

Adding a way to lazy get the function when executing the function may be better. Similar to Flask flask.request that get populated when the function is executed.

Use case / motivation

Use the Airflow context in arbitrary function while keeping the signature of the function stable and easy to reason about.

An idea of implementation would be:

from airflow.task import context

@task 
def context_task(...)
    return context['run_date']

This can be populated by enabling a context manager on the PythonOperator.execute function that populates the field (similar to how DagContext works) and removes it afterwards.

[kaxil]

Just a comment, provide_context is no longer requires in Master: https://github.com/apache/airflow/pull/5990/files/a5d99ed3f3c7b1a815bd2617d3b935bed499f3f7

[casassg]

I see, I still think we should implement this lazy resolved context. Using the current approach from master means that signatures can mismatch op_args/op_kwargs and would make it difficult to reason about the function signature:

@task 
def simple_task(text:str, execution_date: datetime):
  print(text)
  print(execution_date)

simple_task("hello")

As you can see this will execute correctly even though the function signature has a an extra argument. This leads to confusion in my personal opinion hence my proposed solution.

[kaxil]

I see, I still think we should implement this lazy resolved context. Using the current approach from master means that signatures can mismatch op_args/op_kwargs and would make it difficult to reason about the function signature:

@task 
def simple_task(text:str, execution_date: datetime):
  print(text)
  print(execution_date)

simple_task("hello")

As you can see this will execute correctly even though the function signature has a an extra argument. This leads to confusion in my personal opinion hence my proposed solution.

Yup agree, just wanted to make it clear that we don't need provide_context anymore

[jonathanshir]

My suggested changes:
The issue states we want to implement context similarly to flask.requests :

from airflow.task import context
@task 
def context_task(...)
    return context['run_date']

I think this is a bit of an over-complication - our use case is much simpler than this.
We simply have a single execute context for each task instance, so my intended implementation will be a bit simpler:

from airflow.current import current_context
@task
def context_task(...)
    context = current_context()
    return context['run_date']

The actual implementation will contain a simple stack data type that is updated before op.execute is called.
Thoughts?

[turbaszek]

We simply have a single execute context for each task instance, so my intended implementation will be a bit simpler:

from airflow.current import current_context
@task
def context_task(...)
    context = current_context()
    return context['run_date']

+1 for starting simple, what would you say to make this context a namedtuple? Referencing values is much more easier with namedtuple than with dictionary (no need to remember names).

[jonathanshir]

Sounds great, but does this imply we also change the context being provided the old way?

@provide_context
def task(**context):
    print(type(context)) # == NamedTuple?

[turbaszek]

Sounds great, but does this imply we also change the context being provided the old way?

No, if someone uses **context then it implicitly suggest that it will be a dict. I would do this only for the "new" context. Also, we can alter namedtuple __getitem__ to support old (dict like) behaviour.

[casassg]

Saw it on the PR, lgtm as well on making it a function 😄

[casassg]

Any updates on this?