Airflow - find dag run of specific dag id by execution date without time

Question

I would like to find all the dag runs for a specific dag for a specific execution date.

As I read on the documentation there is this function:

dag_runs = DagRun.find(dag_id=self.dag_name, execution_date=datetime.now())

The problem with this is that the time needs to be exactly the same as well. Is there any way that i can pass only the date and will retrieve all the runs no matter what time were during the day?

I know I can filter afterwards, from dag_runs, with a for loop all dags that where in the desired day but I would like something more efficient that does not bring from the db all records.

Using airflow 1.10.10, in gcp composer. So adding a method in the class DagRun is not an option for me.

user_14624409 · Answer

For Airflow >= 2.0.0 you can use:

dag_runs = DagRun.find(
    dag_id=your_dag_id,
    execution_start_date=your_start_date
    execution_end_date=your_end_date
)

For Airflow < 2.0.0 It's possible to create MyDagRun that inherits from DagRun and backport the required functionality.

Here is a working tested code:

from datetime import datetime
from typing import List, Optional, Union

from airflow import DAG
from airflow.models.dagrun import DagRun
from airflow.operators.python_operator import PythonOperator
from airflow.utils import timezone
from airflow.utils.db import provide_session
from sqlalchemy.orm.session import Session


class MyDagRun(DagRun):

    @staticmethod
    @provide_session
    def find(
        dag_id: Optional[Union[str, List[str]]] = None,
        run_id: Optional[str] = None,
        execution_date: Optional[datetime] = None,
        state: Optional[str] = None,
        external_trigger: Optional[bool] = None,
        no_backfills: bool = False,
        session: Session = None,
        execution_start_date: Optional[datetime] = None,
        execution_end_date: Optional[datetime] = None,
    ) -> List["DagRun"]:

        DR = MyDagRun

        qry = session.query(DR)
        dag_ids = [dag_id] if isinstance(dag_id, str) else dag_id
        if dag_ids:
            qry = qry.filter(DR.dag_id.in_(dag_ids))
        if run_id:
            qry = qry.filter(DR.run_id == run_id)
        if execution_date:
            if isinstance(execution_date, list):
                qry = qry.filter(DR.execution_date.in_(execution_date))
            else:
                qry = qry.filter(DR.execution_date == execution_date)
        if execution_start_date and execution_end_date:
            qry = qry.filter(DR.execution_date.between(execution_start_date, execution_end_date))
        elif execution_start_date:
            qry = qry.filter(DR.execution_date >= execution_start_date)
        elif execution_end_date:
            qry = qry.filter(DR.execution_date <= execution_end_date)
        if state:
            qry = qry.filter(DR.state == state)
        if external_trigger is not None:
            qry = qry.filter(DR.external_trigger == external_trigger)
        if no_backfills:
            # in order to prevent a circular dependency
            from airflow.jobs import BackfillJob
            qry = qry.filter(DR.run_id.notlike(BackfillJob.ID_PREFIX + '%'))

        return qry.order_by(DR.execution_date).all()


def func(**kwargs):
    dr = MyDagRun()
    # Need to use timezone to avoid ValueError: naive datetime is disallowed
    start = timezone.make_aware(datetime(2021, 3, 1, 9, 59, 0)) # change to your required start 
    end = timezone.make_aware(datetime(2021, 3, 1, 10, 1, 0)) # change to your required end 
    results = dr.find(execution_start_date=start,
                      execution_end_date=end
                      )

    print("Execution dates met criteria are:")
    for item in results:
        print(item.execution_date)
    return results


default_args = {
    'owner': 'airflow',
    'start_date': datetime(2019, 11, 1),

}

with DAG(dag_id='test',
         default_args=default_args,
         schedule_interval=None,
         catchup=True
         ) as dag:

    op = PythonOperator(task_id="task",
                        python_callable=func)

Example showing 4 existed runs:

Using the code it selected the required runs:

Airflow - find dag run of specific dag id by execution date without time

1 Answers