This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a7d7ef6433 Don't allow defaults other than None in context parameters,
and improve error message (#38015)
a7d7ef6433 is described below
commit a7d7ef6433a068266b1403e1d94dda2dc85bd634
Author: Kevin Languasco <[email protected]>
AuthorDate: Fri Mar 22 16:44:47 2024 -0500
Don't allow defaults other than None in context parameters, and improve
error message (#38015)
* feat: don't allow defaults other than None for context parameters
* feat: improve error message when replacing context parameters with None
breaks signature
* feat: improve error message for unsupported position of context key
parameter
* style: make stack trace more readable by putting the error message in a
variable
* feat: add details to exceptions and related test
* fix: put task decorated function outside pytest.raises call
---
airflow/decorators/base.py | 29 ++++++++++++++++++++++++++++-
tests/decorators/test_python.py | 11 +++++++++++
2 files changed, 39 insertions(+), 1 deletion(-)
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index 0b0d921f8c..7adaf8a447 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -208,11 +208,38 @@ class DecoratedOperator(BaseOperator):
# 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)
+
+ # Don't allow context argument defaults other than None to avoid
ambiguities.
+ faulty_parameters = [
+ param.name
+ for param in signature.parameters.values()
+ if param.name in KNOWN_CONTEXT_KEYS and param.default not in
(None, inspect.Parameter.empty)
+ ]
+ if faulty_parameters:
+ message = f"Context key parameter {faulty_parameters[0]} can't
have a default other than None"
+ raise ValueError(message)
+
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)
+ try:
+ signature = signature.replace(parameters=parameters)
+ except ValueError as err:
+ message = textwrap.dedent(
+ f"""
+ The function signature broke while assigning defaults to
context key parameters.
+
+ The decorator is replacing the signature
+ > {python_callable.__name__}({', '.join(str(param) for param
in signature.parameters.values())})
+
+ with
+ > {python_callable.__name__}({', '.join(str(param) for param
in parameters)})
+
+ which isn't valid: {err}
+ """
+ )
+ raise ValueError(message) from err
# 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
diff --git a/tests/decorators/test_python.py b/tests/decorators/test_python.py
index fde79a4784..5175c508ac 100644
--- a/tests/decorators/test_python.py
+++ b/tests/decorators/test_python.py
@@ -256,6 +256,17 @@ class TestAirflowTaskDecorator(BasePythonTest):
add_number()
add_number("test")
+ def test_fails_context_parameter_other_than_none(self):
+ """Fail if a context parameter has a default and it's not None."""
+ error_message = "Context key parameter try_number can't have a default
other than None"
+
+ @task_decorator
+ def add_number_to_try_number(num: int, try_number: int = 0):
+ return num + try_number
+
+ with pytest.raises(ValueError, match=error_message):
+ add_number_to_try_number(1)
+
def test_fail_method(self):
"""Tests that @task will fail if signature is not binding."""
