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 <kevin+git...@languas.co> 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."""