Airflow Task Dependency Sensor

The AirflowTaskDependencySensor is a native brickflow sensor that polls a remote Airflow cluster's REST API to wait for a specific DAG task to reach an allowed state. It has no dependency on apache-airflow on the Databricks side -- it only needs requests.

Prior versions shipped a TaskDependencySensor and an AirflowProxyOktaClusterAuth class that both subclassed Airflow. Those classes are still importable from brickflow_plugins but now raise RuntimeError on instantiation. Migrate to AirflowTaskDependencySensor and the plain AirflowCluster dataclass:

from datetime import timedelta

from brickflow_plugins import AirflowCluster, AirflowTaskDependencySensor

sensor = AirflowTaskDependencySensor(
    dag_id="my_upstream_dag",
    task_id="final_task",
    cluster=AirflowCluster(
        url="https://airflow.example.com",
        version="2.0.2",
        token=my_bearer_token,  # compute this from Okta/etc. yourself
    ),
    allowed_states=["success"],
    execution_delta=timedelta(hours=0),
    timeout_seconds=3600,
    poke_interval=60,
)
sensor.execute()

Supported Airflow versions

The sensor auto-routes to the correct REST API dialect based on the AirflowCluster.version string:

version starts with	Endpoint prefix	Notes
"1."	/api/experimental	Airflow 1.x
any other value (default)	/api/v1	Airflow 2.x
"3."	/api/v2	Airflow 3.x (FastAPI). Uses logical_date_gte in place of execution_date_gte, and drops asset-triggered runs that have logical_date=null.

Airflow 3.x example

from datetime import timedelta

from brickflow_plugins import AirflowCluster, AirflowTaskDependencySensor

sensor = AirflowTaskDependencySensor(
    dag_id="my_upstream_dag",
    task_id="final_task",
    cluster=AirflowCluster(
        url="https://airflow.example.com",
        version="3.0.0",
        token=my_jwt,  # short-lived JWT from Okta/MAP or POST /auth/token
    ),
    allowed_states=["success"],
    execution_delta=timedelta(hours=0),
    timeout_seconds=3600,
    poke_interval=60,
)
sensor.execute()

Auth is unchanged: bring your own bearer token in cluster.token. The sensor does not call /auth/token for you.

API Reference

Airflow Task Dependency Sensor.

Native brickflow sensor that polls an external Airflow API to check the status of a specific task in an Airflow DAG. Supports Airflow 1.x (/api/experimental), 2.x (/api/v1), and 3.x (/api/v2) API shapes. Requires only requests -- no apache-airflow package needs to be installed on the Databricks cluster.

Classes

brickflow_plugins.sensors.airflow_task_dependency_sensor.AirflowCluster(url: str, version: str, token: str)

Represents an Airflow cluster the sensor polls.

Parameters

url : str Base URL of the Airflow API (e.g. https://airflow.example.com). version : str Airflow major version string, e.g. "1.10", "2.0.2", or "3.0.0". Used to select between the /api/experimental (Airflow 1.x), /api/v1 (Airflow 2.x), and /api/v2 (Airflow 3.x) endpoint shapes. token : str Bearer token that will be sent in the Authorization header. For Airflow 3.x, this is typically a short-lived JWT obtained out-of-band (e.g. via Okta, MAP, or POST /auth/token).

Source code in brickflow_plugins/sensors/airflow_task_dependency_sensor.py

def __init__(self, url: str, version: str, token: str) -> None:
    self.url = str(url).rstrip("/")
    self.version = version
    self.token = token

Attributes

token = token instance-attribute

url = str(url).rstrip('/') instance-attribute

version = version instance-attribute

brickflow_plugins.sensors.airflow_task_dependency_sensor.AirflowTaskDependencySensor(dag_id: str, task_id: str, cluster: AirflowCluster, allowed_states: Optional[List[str]] = None, execution_delta: timedelta = timedelta(days=0), latest: bool = False, timeout_seconds: int = 3600, poke_interval: int = 60)

Bases: Sensor

Sensor that polls an external Airflow cluster's API to wait until a given task in a given DAG reaches an allowed state.

The API dialect used is selected from cluster.version:

• "1.x" -> /api/experimental
• "2.x" (default) -> /api/v1
• "3.x" -> /api/v2 (FastAPI, logical_date filters)

Example

::

sensor = AirflowTaskDependencySensor(
    dag_id="my_upstream_dag",
    task_id="final_task",
    cluster=AirflowCluster(
        url="https://airflow.example.com",
        version="2.0.2",   # use "3.0.0" for Airflow 3.x (/api/v2)
        token=my_token,
    ),
    execution_delta=timedelta(hours=0),
    timeout_seconds=3600,
    poke_interval=60,
)
sensor.execute()

Source code in brickflow_plugins/sensors/airflow_task_dependency_sensor.py

def __init__(
    self,
    dag_id: str,
    task_id: str,
    cluster: AirflowCluster,
    allowed_states: Optional[List[str]] = None,
    execution_delta: timedelta = timedelta(days=0),
    latest: bool = False,
    timeout_seconds: int = 3600,
    poke_interval: int = 60,
) -> None:
    super().__init__()
    self.dag_id = dag_id
    self.task_id = task_id
    self.cluster = cluster
    self.allowed_states = allowed_states if allowed_states else ["success"]
    self.execution_delta = execution_delta
    self.latest = latest
    self.poke_interval = poke_interval
    self.timeout = timeout_seconds

    self._poke_count = 0
    self._start_time = time.time()

Attributes

allowed_states = allowed_states if allowed_states else ['success'] instance-attribute

cluster = cluster instance-attribute

dag_id = dag_id instance-attribute

execution_delta = execution_delta instance-attribute

latest = latest instance-attribute

poke_interval = poke_interval instance-attribute

task_id = task_id instance-attribute

timeout = timeout_seconds instance-attribute

Functions

execute() -> None

Poll the Airflow API until the task reaches an allowed state or the timeout is exceeded.

Raises

TimeoutError If self.timeout seconds elapse before the task reaches an allowed state.

Source code in brickflow_plugins/sensors/airflow_task_dependency_sensor.py

def execute(self) -> None:
    """
    Poll the Airflow API until the task reaches an allowed state or the
    timeout is exceeded.

    Raises
    ------
    TimeoutError
        If ``self.timeout`` seconds elapse before the task reaches an
        allowed state.
    """
    log.info("Execution date derived from context: %s", self._execution_timestamp)

    execution_window_tz = self._execution_timestamp + self.execution_delta
    log.info(
        "Executing the sensor to check for %s for %s DAG and task %s after %s.",
        self.allowed_states,
        self.dag_id,
        self.task_id,
        execution_window_tz,
    )
    status = ""
    while status not in self.allowed_states:
        status = self.poke()
        if status == "failed":
            # Log the fact that upstream failed, however do not fail the task
            # and continue poking until timeout.
            log.error(
                "Upstream dag '%s' failed at '%s' task, continue poking till "
                "timeout is reached...",
                self.dag_id,
                self.task_id,
            )
            time.sleep(self.poke_interval)
        elif status != "success":
            time.sleep(self.poke_interval)

        if (time.time() - self._start_time) > self.timeout:
            raise TimeoutError("The job has timed out!")
    log.info("Upstream DAG '%s' is successful", self.dag_id)

get_execution_stats(execution_date: datetime, max_end_date: Optional[datetime] = None) -> str

Return the state of self.task_id for the most recent self.dag_id DAG run in the given window.

Returns "none" when no matching DAG run is found.

Source code in brickflow_plugins/sensors/airflow_task_dependency_sensor.py

def get_execution_stats(
    self,
    execution_date: datetime,
    max_end_date: Optional[datetime] = None,
) -> str:
    """
    Return the state of ``self.task_id`` for the most recent
    ``self.dag_id`` DAG run in the given window.

    Returns "none" when no matching DAG run is found.
    """
    variant = _api_variant(self.cluster.version)
    execution_window_tz = (execution_date + self.execution_delta).strftime(
        "%Y-%m-%dT%H:%M:%SZ"
    )
    max_end_date_filter = (
        f"&end_date_lte={max_end_date.strftime('%Y-%m-%dT%H:%M:%SZ')}"
        if max_end_date
        else ""
    )
    headers = {
        "Content-Type": "application/json",
        "cache-control": "no-cache",
        "Authorization": f"Bearer {self.cluster.token}",
    }
    if variant == "experimental":
        log.info("this is 1.x cluster")
        url = f"{self.cluster.url}/api/experimental/dags/{self.dag_id}/dag_runs/"
    elif variant == "v2":
        # Airflow 3.x FastAPI: /api/v2 replaces /api/v1, and the
        # execution_date_* filters were replaced by logical_date_*.
        url = (
            f"{self.cluster.url}/api/v2/dags/{self.dag_id}"
            f"/dagRuns?logical_date_gte={execution_window_tz}{max_end_date_filter}"
        )
    else:
        # Airflow 2.x API limits 100 records, so only picking runs
        # within the execution window provided.
        url = (
            f"{self.cluster.url}/api/v1/dags/{self.dag_id}"
            f"/dagRuns?execution_date_gte={execution_window_tz}{max_end_date_filter}"
        )

    log.info("URL to poke for dag runs %s", url)
    response = requests.get(url, headers=headers, verify=False, timeout=10)
    response.raise_for_status()

    list_of_dictionaries = response.json()["dag_runs"]
    # Airflow 3.x drops `execution_date` from DagRun payloads in favor of
    # `logical_date`; older APIs still expose `execution_date`.
    sort_key = "logical_date" if variant == "v2" else "execution_date"
    if variant == "v2":
        # In Airflow 3.x `logical_date` is nullable for asset-triggered runs.
        # This sensor is fundamentally a date-window check, so runs without
        # a logical_date are not eligible dependency targets -- drop them
        # before sorting so they can't be selected by `[-1]` / `[0]`.
        list_of_dictionaries = [
            r for r in list_of_dictionaries if r.get(sort_key) is not None
        ]
    list_of_dictionaries = sorted(
        list_of_dictionaries,
        key=lambda k: k[sort_key],
        reverse=True,
    )

    if len(list_of_dictionaries) == 0:
        log.info(
            "No runs found for %s dag in time window: %s - %s, please check "
            "upstream dag",
            self.dag_id,
            execution_window_tz,
            max_end_date.strftime("%Y-%m-%dT%H:%M:%SZ") if max_end_date else "now",
        )
        return "none"

    if variant == "experimental":
        # For Airflow 1.x the execution date is needed to check the status.
        dag_run_id = list_of_dictionaries[0]["execution_date"]
    else:
        # For Airflow 2.x/3.x the dag_run_id is needed to check the status.
        dag_run_id = (
            list_of_dictionaries[-1]["dag_run_id"]
            if not self.latest
            else list_of_dictionaries[0]["dag_run_id"]
        )

    log.info("Latest run for the dag is with execution date of %s", dag_run_id)
    log.info(
        "Poking %s dag for %s run_id status as latest flag is set to %s",
        self.dag_id,
        dag_run_id,
        self.latest,
    )

    if variant == "experimental":
        if dag_run_id >= execution_window_tz:
            task_url = f"{url}/{dag_run_id}/tasks/{self.task_id}"
        else:
            log.info(
                "No airflow runs found for %s dag after %s",
                self.dag_id,
                execution_window_tz,
            )
            return "none"
    else:
        api_prefix = "/api/v2" if variant == "v2" else "/api/v1"
        task_url = (
            f"{self.cluster.url}{api_prefix}/dags/{self.dag_id}"
            f"/dagRuns/{dag_run_id}/taskInstances/{self.task_id}"
        )
    log.info("Pinging airflow API %s for task status ", task_url)
    task_response = requests.get(
        task_url, headers=headers, verify=False, timeout=10
    )
    task_response.raise_for_status()
    return task_response.json()["state"]

poke() -> str

Poke the Airflow API once and return the task state.

Source code in brickflow_plugins/sensors/airflow_task_dependency_sensor.py

def poke(self) -> str:  # type: ignore[override]
    """Poke the Airflow API once and return the task state."""
    log.info("executing poke... %s", self._poke_count)
    self._poke_count += 1
    log.info("Poking... %s round", self._poke_count)

    task_status = self.get_execution_stats(execution_date=self._execution_timestamp)
    log.info("task_status=%s", task_status)
    return task_status

Autosys Sensor

Autosys Sensor.

Native brickflow sensor that polls the Autosys REST API to wait for a job to reach a successful state. Requires only requests -- no apache-airflow package needs to be installed on the Databricks cluster.

Classes

brickflow_plugins.sensors.autosys_sensor.AutosysSensor(url: str, job_name: str, poke_interval: int, time_delta: Union[timedelta, dict] = timedelta(days=0))

Bases: Sensor

Sensor that polls an Autosys REST endpoint for the given job_name and waits until it reports a successful status recent enough to satisfy time_delta.

Example

::

sensor = AutosysSensor(
    url="https://autosys.example.com/api/jobs",
    job_name="my_upstream_job",
    poke_interval=60,
    time_delta=timedelta(hours=1),
)
sensor.poke()

Source code in brickflow_plugins/sensors/autosys_sensor.py

def __init__(
    self,
    url: str,
    job_name: str,
    poke_interval: int,
    time_delta: Union[timedelta, dict] = timedelta(days=0),
) -> None:
    super().__init__()
    self.url = str(url).rstrip("/")
    self.job_name = job_name
    self.poke_interval = poke_interval
    self.time_delta = time_delta

Attributes

job_name = job_name instance-attribute

poke_interval = poke_interval instance-attribute

time_delta = time_delta instance-attribute

url = str(url).rstrip('/') instance-attribute

Functions

poke()

Poke the Autosys API once. Recurses (via time.sleep + self-call) until success.

Source code in brickflow_plugins/sensors/autosys_sensor.py

def poke(self):  # type: ignore[override]
    """Poke the Autosys API once. Recurses (via ``time.sleep`` + self-call) until success."""
    url = f"{self.url}/{self.job_name}"
    log.info("Poking: %s", url)

    headers = {
        "Accept": "application/json",
        "cache-control": "no-cache",
    }

    response = requests.get(
        url,
        headers=headers,
        verify=False,  # nosec
        timeout=10,
    )

    if response.status_code != 200:
        raise HTTPError(
            f"Request failed with '{response.status_code}' code. \n{response.text}"
        )

    status = response.json()["status"][:2].upper()

    last_end_timestamp = None
    if last_end_utc := response.json().get("lastEndUTC"):
        last_end_timestamp = parse(last_end_utc).replace(tzinfo=pytz.UTC)

    time_delta = (
        self.time_delta
        if isinstance(self.time_delta, timedelta)
        else timedelta(**self.time_delta)
    )

    run_timestamp = self._execution_timestamp - time_delta

    if (
        "SU" in status
        and last_end_timestamp
        and last_end_timestamp >= run_timestamp
    ):
        log.info(
            "Last End: %s, Run Timestamp: %s", last_end_timestamp, run_timestamp
        )
        log.info("Success criteria met. Exiting")
        return True

    log.info("Last End: %s, Run Timestamp: %s", last_end_timestamp, run_timestamp)
    time.sleep(self.poke_interval)
    log.info("Poking again")
    return AutosysSensor.poke(self)