ExternalTaskSensor (only for last dag runs) in Airflow 3

[Frost3281]

In Airflow 2, I used the following Sensor to check the status of the latest run of external DAGs. In Airflow 3, this code no longer works.

"Direct database access via the ORM is not allowed in Airflow 3.0"

I have tried many options: Engine / Connection and other methods of database access. And I tried access through API (always getting connection refused). Is there a proper way in Airflow 3 to achieve the same functionality with this operator?

import json
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, cast

import pytz  # type: ignore
from airflow.exceptions import TaskNotFound
from airflow.models import DagBag, DagModel
from airflow.sensors.external_task import ExternalTaskSensor
from airflow.utils.db import provide_session  # type: ignore
from sqlalchemy import select
from sqlalchemy.orm import Session

DAGS_TO_EXTERNAL_TASK_DOCS: dict[str, str] = {}


class LastDagRunSensor(ExternalTaskSensor):
    """Sensor to retrieve the last run of a DAG."""

    def __init__(
        self,
        external_dag_id: str,
        external_task_id: str | None = None,
        **kwargs: Any,
    ) -> None:
        """Initializer."""

        def dag_last_exec(crnt_dttm: datetime) -> datetime:  # noqa: ARG001
            return self.get_dag_last_execution_date(self.external_dag_id)  # type: ignore

        super().__init__(
            external_dag_id=external_dag_id,
            external_task_id=external_task_id,
            **kwargs,
        )
        self.execution_date_fn = dag_last_exec

    @provide_session
    def get_dag_last_execution_date(
        self,
        dag_id: str,  # noqa: ARG002
        session: Session,
    ) -> datetime:
        """Get the last execution date of the DAG."""
        dag = _get_dag_by_id(session, self.external_dag_id)
        if not dag:
            return _yesterday()
        return dag.get_last_dagrun(include_externally_triggered=True).execution_date


def _get_dag_by_id(session: Session, dag_id: str) -> DagModel | None:
    query = select(DagModel).where(DagModel.dag_id == dag_id)
    return session.execute(query).scalars().first()


def _yesterday() -> datetime:
    return datetime.now(tz=pytz.timezone('Europe/Moscow')) - timedelta(days=1)

[potiuk]

I have tried many options: Engine / Connection and other methods of database access. And I tried access through API (always getting connection refused). Is there a proper way in Airflow 3 to achieve the same functionality with this operator?

Nope. But if you get "connection refused" it likely means you are using wrong hostname, wrong port, or you have to configure your networking to allow your workers to connect to your api_server (using the right host and port).

Connection refused always means tha when you try to connect to a give host and port, there is nothing on the other side to listen to your connection - so It's either firewall/networking blocking you, or you use wrong host/port.

[potiuk]

Security trumps convenience, sorry. There is no way the fact that you can write few less line of code makes it an argument to be able to drop the whole database. The impact of those two things is incomparable.

And with REST API and python client of ours - you do not even have to make it very complex. Python client exposes Pythonic interface that you can use to query for stuff. Being a little sarcastic - I know it;s an awfully lot of work to install one more package, configure a user that can be used via REST API and run maybe 10 lines of code to do the same things with python client and REST API. Sorry for making you to do all that work, I know it's really difficult, but in exchange you get much more secure setup of Airflow.

You are welcome.

[ajbosco]

As far as I can tell, the API does not expose the ability to do something like this:

 if self.lookback_date == "run_date":
            count = (
                session.query(TI)
                .filter(
                    TI.dag_id == self.external_dag_id,
                    TI.task_id == self.external_task_id,
                    TI.state.in_(self.allowed_states),
                    TI.start_date >= start_dttm,
                    TI.start_date <= end_dttm,
                )
                .count()
            )
        else:
            count = (
                session.query(TI)
                .filter(
                    TI.dag_id == self.external_dag_id,
                    TI.task_id == self.external_task_id,
                    TI.state.in_(self.allowed_states),
                    TI.logical_date >= start_dttm,
                    TI.logical_date <= end_dttm,
                )
                .count()
            )

[gopidesupavan]

There are already functionality implemeted for getting count and states etc; look at this
https://github.com/apache/airflow/blob/main/providers/standard/src/airflow/providers/standard/triggers/external_task.py#L142

from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
RuntimeTaskInstance.get_ti_count
RuntimeTaskInstance.get_task_states
RuntimeTaskInstance.get_dr_count
RuntimeTaskInstance.get_dagrun_state
RuntimeTaskInstance.get_dr_count

You can take a look and see whats useful in your scenario for above

[ajbosco]

I tried get_ti_count, but it takes a list of logical_dates not a start/end date.

[potiuk]

As far as I can tell, the API does not expose the ability to do something like this:

@ajbosco - I am sure you can run appropriate python code to query REST API - did you actually try it?

And if you think there is a use case which might be also useful for others - to be "easily" accessible via Task SDK API - as explained before - please make a feature request, describing the use case and why you think it's good to have it as a "feature" in the Task SDKL - it's all in your hands to describe it well enough and justify that this case is useful for many users. Because we are optimising for many users. It's a little sad to hear but if there is a case that useful for you only- maintainers who give you the software for free - often working in their evenings and weekends will not optimise the acess for "you" only, That would be quite selfish from your side if you think it's justified to make it easy only for you.

But if you justify that it is useful for others as well, and nicely describe it in a request, there are already things we added betwen 3.0.0 and 3.0.4 that we got convinced are useful for many users and added.

But it's up to you to describe it in the way that will be convincing. All in your hands.

[gopidesupavan]

@Frost3281 you cant access db directly from workers.