Fix access to context in functions decorated by task.sensor

[Thomas-McKanna]

Description

Hello, I am a new Airflow user. I am requesting a feature in which the airflow context (containing task instance, etc.) be available inside of functions decorated by airflow.decorators.task.sensor.

Use case/motivation

I have noticed that when using the airflow.decorators.task decorator, one can access items from the context (such as the task instance) by using **kwargs or keyword arguments in the decorated function. But I have discovered that the same is not true for the airflow.decorators.task.sensor decorator. I'm not sure if this is a bug or intentional, but it would be very useful to be able to access the context normally from functions decorated by task.sensor.

I believe this may have been an oversight. The DecoratedSensorOperator class is a child class of PythonSensor:

airflow/airflow/decorators/sensor.py

Line 28 in 1fbfd31

class DecoratedSensorOperator(PythonSensor):

This DecoratedSensorOperator class overrides poke, but does not incorporate the passed in Context object before calling the decorated function:

airflow/airflow/decorators/sensor.py

Lines 60 to 61 in 1fbfd31

	def poke(self, context: Context) -> PokeReturnValue | bool:
	return self.python_callable(*self.op_args, **self.op_kwargs)

This is in contrast to the PythonSensor, whose poke method merges the context with the existing op_kwargs:

airflow/airflow/sensors/python.py

Lines 68 to 77 in 1fbfd31

	def poke(self, context: Context) -> PokeReturnValue | bool:
	context_merge(context, self.op_kwargs, templates_dict=self.templates_dict)
	self.op_kwargs = determine_kwargs(self.python_callable, self.op_args, context)
	self.log.info("Poking callable: %s", str(self.python_callable))
	return_value = self.python_callable(*self.op_args, **self.op_kwargs)
	if isinstance(return_value, PokeReturnValue):
	return return_value
	else:
	return PokeReturnValue(bool(return_value))

This seems like an easy fix, and I'd be happy to submit a pull request. But I figured I'd start with a feature request since I'm new to the open source community.

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template!

[Taragolis]

Feel free to make a PR

[BrunoGrandePhD]

I see that a PR to address this is already open and already approved by at least one reviewer, but if anyone else stumbles upon this issue, you can use airflow.operators.python.get_current_context() as a temporary solution to access context variables (docs).