Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow context key args to not provide a default #33430

Merged
merged 1 commit into from Aug 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
15 changes: 13 additions & 2 deletions airflow/decorators/base.py
Expand Up @@ -196,16 +196,27 @@ def __init__(
op_args = op_args or []
op_kwargs = op_kwargs or {}

# Check the decorated function's signature. We go through the argument
# list and "fill in" defaults to arguments that are known context keys,
# since values for those will be provided when the task is run. Since
# we're not actually running the function, None is good enough here.
signature = inspect.signature(python_callable)
parameters = [
param.replace(default=None) if param.name in KNOWN_CONTEXT_KEYS else param
for param in signature.parameters.values()
]
signature = signature.replace(parameters=parameters)

# Check that arguments can be binded. There's a slight difference when
# we do validation for task-mapping: Since there's no guarantee we can
# receive enough arguments at parse time, we use bind_partial to simply
# check all the arguments we know are valid. Whether these are enough
# can only be known at execution time, when unmapping happens, and this
# is called without the _airflow_mapped_validation_only flag.
if kwargs.get("_airflow_mapped_validation_only"):
inspect.signature(python_callable).bind_partial(*op_args, **op_kwargs)
signature.bind_partial(*op_args, **op_kwargs)
else:
inspect.signature(python_callable).bind(*op_args, **op_kwargs)
signature.bind(*op_args, **op_kwargs)

self.multiple_outputs = multiple_outputs
self.op_args = op_args
Expand Down
26 changes: 16 additions & 10 deletions docs/apache-airflow/tutorial/taskflow.rst
Expand Up @@ -578,23 +578,29 @@ task to copy the same file to a date-partitioned storage location in S3 for long
Accessing context variables in decorated tasks
----------------------------------------------

When running your callable, Airflow will pass a set of keyword arguments that can be used in your
function. This set of kwargs correspond exactly to what you can use in your Jinja templates.
For this to work, you need to define ``**kwargs`` in your function header, or you can add directly the
keyword arguments you would like to get - for example with the below code your callable will get
the values of ``ti`` and ``next_ds`` context variables. Note that when explicit keyword arguments are used,
they must be made optional in the function header to avoid ``TypeError`` exceptions during DAG parsing as
these values are not available until task execution.
When running your callable, Airflow will pass a set of keyword arguments that
can be used in your function. This set of kwargs correspond exactly to what you
can use in your Jinja templates. For this to work, you can add context keys you
would like to receive in the function as keyword arguments.

With explicit arguments:
For example, the callable in the code block below will get values of the ``ti``
and ``next_ds`` context variables:

.. code-block:: python

@task
def my_python_callable(ti=None, next_ds=None):
def my_python_callable(*, ti, next_ds):
pass

With kwargs:
.. versionchanged:: 2.8
Previously the context key arguments must provide a default, e.g. ``ti=None``.
This is no longer needed.

You can also choose to receive the entire context with ``**kwargs``. Note that
this can incur a slight performance penalty since Airflow will need to
expand the entire context that likely contains many things you don't actually
need. It is therefore more recommended for you to use explicit arguments, as
demonstrated in the previous paragraph.

.. code-block:: python

Expand Down