In this lesson, you will learn how to extend the base BaseSensorOperator to create specialized logic for waiting on non-standard signals that off-the-shelf sensors cannot handle.
While built-in sensors like FileSensor or HttpSensor cover common scenarios, real-world data pipelines often encounter "unique wait conditions." Perhaps you need to monitor a specific log entry in a legacy system, wait for a hardware-level status flag via a custom API, or check a database table for a very specific state transition that isn't a simple count.
To build a custom sensor, you must inherit from BaseSensorOperator and implement the poke method. This method should return True when your condition is met and False otherwise. If the method returns False, the sensor will pause execution and check again later, based on your configured poke_interval.
Here is a simplified example of a sensor waiting for a custom API response status:
from airflow.sensors.base import BaseSensorOperator
from my_custom_hook import MyCustomHook
class CustomStatusSensor(BaseSensorOperator):
def __init__(self, target_status, *args, **kwargs):
super().__init__(*args, **kwargs)
self.target_status = target_status
self.hook = MyCustomHook()
def poke(self, context):
current_status = self.hook.get_status()
self.log.info(f"Checking status... Current: {current_status}")
# Return True to stop the sensor, False to continue polling
return current_status == self.target_status
By encapsulating the connection logic inside a Hook (which we discussed in previous lessons) and the specific wait condition within the poke method, you create a reusable, modular component. This approach keeps your DAGs clean, as you can simply instantiate CustomStatusSensor in your workflow rather than embedding complex retry logic or polling loops inside a generic PythonOperator.
When designing these, ensure your poke method is lightweight. Because sensors occupy a worker slot, perform only the bare minimum API or file system calls required to evaluate the state. If the wait condition is expected to take hours, remember to utilize mode='reschedule' to release the worker slot between pokes, keeping your resource allocation efficient.
Key Takeaway: Build custom sensors by inheriting from BaseSensorOperator and defining a precise poke method to encapsulate domain-specific wait logic, ensuring your orchestration layer remains clean and maintainable.