Create a more efficient airflow dag test command that also has better local logging

airflow/cli/cli_parser.py

@@ -2035,7 +2035,7 @@ def _remove_dag_id_opt(command: ActionCommand):
         subcommands=[
             _remove_dag_id_opt(sp)
             for sp in DAGS_COMMANDS
-            if sp.name in ['backfill', 'list-runs', 'pause', 'unpause']
+            if sp.name in ['backfill', 'list-runs', 'pause', 'unpause', 'test']
         ],
     ),
     GroupCommand(

airflow/cli/commands/dag_command.py

@@ -33,8 +33,7 @@
 from airflow.api.client import get_current_api_client
 from airflow.cli.simple_table import AirflowConsole
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, BackfillUnfinished, RemovedInAirflow3Warning
-from airflow.executors.debug_executor import DebugExecutor
+from airflow.exceptions import AirflowException, RemovedInAirflow3Warning
 from airflow.jobs.base_job import BaseJob
 from airflow.models import DagBag, DagModel, DagRun, TaskInstance
 from airflow.models.dag import DAG
@@ -449,7 +448,7 @@ def dag_list_dag_runs(args, dag=None, session=NEW_SESSION):
 
 @provide_session
 @cli_utils.action_cli
-def dag_test(args, session=None):
+def dag_test(args, dag=None, session=None):
     """Execute one single DagRun for a given DAG and execution date, using the DebugExecutor."""
     run_conf = None
     if args.conf:
@@ -458,22 +457,8 @@ def dag_test(args, session=None):
         except ValueError as e:
             raise SystemExit(f"Configuration {args.conf!r} is not valid JSON. Error: {e}")
     execution_date = args.execution_date or timezone.utcnow()
-    dag = get_dag(subdir=args.subdir, dag_id=args.dag_id)
-    dag.clear(start_date=execution_date, end_date=execution_date, dag_run_state=False)
-    try:
-        dag.run(
-            executor=DebugExecutor(),
-            start_date=execution_date,
-            end_date=execution_date,
-            conf=run_conf,
-            # Always run the DAG at least once even if no logical runs are
-            # available. This does not make a lot of sense, but Airflow has
-            # been doing this prior to 2.2 so we keep compatibility.
-            run_at_least_once=True,
-        )
-    except BackfillUnfinished as e:
-        print(str(e))
-
+    dag = dag or get_dag(subdir=args.subdir, dag_id=args.dag_id)
+    dag.test(execution_date=execution_date, run_conf=run_conf, session=session)
     show_dagrun = args.show_dagrun
     imgcat = args.imgcat_dagrun
     filename = args.save_dagrun

airflow/example_dags/example_bash_operator.py

@@ -73,4 +73,4 @@
 this_will_skip >> run_this_last
 
 if __name__ == "__main__":
-    dag.cli()
+    dag.test()

airflow/models/dag.py

@@ -63,10 +63,11 @@
 import airflow.templates
 from airflow import settings, utils
 from airflow.compat.functools import cached_property
-from airflow.configuration import conf
+from airflow.configuration import conf, secrets_backend_list
 from airflow.exceptions import (
     AirflowDagInconsistent,
     AirflowException,
+    AirflowSkipException,
     DuplicateTaskIdFound,
     RemovedInAirflow3Warning,
     TaskNotFound,
@@ -80,6 +81,7 @@
 from airflow.models.operator import Operator
 from airflow.models.param import DagParam, ParamsDict
 from airflow.models.taskinstance import Context, TaskInstance, TaskInstanceKey, clear_task_instances
+from airflow.secrets.local_filesystem import LocalFilesystemBackend
 from airflow.security import permissions
 from airflow.stats import Stats
 from airflow.timetables.base import DagRunInfo, DataInterval, TimeRestriction, Timetable
@@ -2435,6 +2437,74 @@ def cli(self):
         args = parser.parse_args()
         args.func(args, self)
 
+    @provide_session
+    def test(
+        self,
+        execution_date: datetime | None = None,
+        run_conf: dict[str, Any] | None = None,
+        conn_file_path: str | None = None,
+        variable_file_path: str | None = None,
+        session: Session = NEW_SESSION,
+    ) -> None:
+        """Execute one single DagRun for a given DAG and execution date."""
+
+        def add_logger_if_needed(ti: TaskInstance):
+            """
+            Add a formatted logger to the taskinstance so all logs are surfaced to the command line instead
+            of into a task file. Since this is a local test run, it is much better for the user to see logs
+            in the command line, rather than needing to search for a log file.
+            Args:
+                ti: The taskinstance that will receive a logger
+
+            """
+            format = logging.Formatter("[%(asctime)s] {%(filename)s:%(lineno)d} %(levelname)s - %(message)s")
+            handler = logging.StreamHandler(sys.stdout)
+            handler.level = logging.INFO
+            handler.setFormatter(format)
+            # only add log handler once
+            if not any(isinstance(h, logging.StreamHandler) for h in ti.log.handlers):
+                self.log.debug("Adding Streamhandler to taskinstance %s", ti.task_id)
+                ti.log.addHandler(handler)
+
+        if conn_file_path or variable_file_path:
+            local_secrets = LocalFilesystemBackend(
+                variables_file_path=variable_file_path, connections_file_path=conn_file_path
+            )
+            secrets_backend_list.insert(0, local_secrets)
+
+        execution_date = execution_date or timezone.utcnow()
+        self.log.debug("Clearing existing task instances for execution date %s", execution_date)
+        self.clear(
+            start_date=execution_date,
+            end_date=execution_date,
+            dag_run_state=False,  # type: ignore
+            session=session,
+        )
+        self.log.debug("Getting dagrun for dag %s", self.dag_id)
+        dr: DagRun = _get_or_create_dagrun(
+            dag=self,
+            start_date=execution_date,
+            execution_date=execution_date,
+            run_id=DagRun.generate_run_id(DagRunType.MANUAL, execution_date),
+            session=session,
+            conf=run_conf,
+        )
+
+        tasks = self.task_dict
+        self.log.debug("starting dagrun")
+        # Instead of starting a scheduler, we run the minimal loop possible to check
+        # for task readiness and dependency management. This is notably faster
+        # than creating a BackfillJob and allows us to surface logs to the user
+        while dr.state == State.RUNNING:
+            schedulable_tis, _ = dr.update_state(session=session)
+            for ti in schedulable_tis:
+                add_logger_if_needed(ti)
+                ti.task = tasks[ti.task_id]
+                _run_task(ti, session=session)
+        if conn_file_path or variable_file_path:
+            # Remove the local variables we have added to the secrets_backend_list
+            secrets_backend_list.pop(0)
+
     @provide_session
     def create_dagrun(
         self,
@@ -3505,3 +3575,66 @@ def get_current_dag(cls) -> DAG | None:
             return cls._context_managed_dags[0]
         except IndexError:
             return None
+
+
+def _run_task(ti: TaskInstance, session):
+    """
+    Run a single task instance, and push result to Xcom for downstream tasks. Bypasses a lot of
+    extra steps used in `task.run` to keep our local running as fast as possible
+    This function is only meant for the `dag.test` function as a helper function.
+
+    Args:
+        ti: TaskInstance to run
+    """
+    log.info("*****************************************************")
+    if ti.map_index > 0:
+        log.info("Running task %s index %d", ti.task_id, ti.map_index)
+    else:
+        log.info("Running task %s", ti.task_id)
+    try:
+        ti._run_raw_task(session=session)
+        session.flush()
+        log.info("%s ran successfully!", ti.task_id)
+    except AirflowSkipException:
+        log.info("Task Skipped, continuing")
+    log.info("*****************************************************")
+
+
+def _get_or_create_dagrun(
+    dag: DAG,
+    conf: dict[Any, Any] | None,
+    start_date: datetime,
+    execution_date: datetime,
+    run_id: str,
+    session: Session,
+) -> DagRun:
+    """
+    Create a DAGRun, but only after clearing the previous instance of said dagrun to prevent collisions.
+    This function is only meant for the `dag.test` function as a helper function.
+    :param dag: Dag to be used to find dagrun
+    :param conf: configuration to pass to newly created dagrun
+    :param start_date: start date of new dagrun, defaults to execution_date
+    :param execution_date: execution_date for finding the dagrun
+    :param run_id: run_id to pass to new dagrun
+    :param session: sqlalchemy session
+    :return:
+    """
+    log.info("dagrun id: %s", dag.dag_id)
+    dr: DagRun = (
+        session.query(DagRun)
+        .filter(DagRun.dag_id == dag.dag_id, DagRun.execution_date == execution_date)
+        .first()
+    )
+    if dr:
+        session.delete(dr)
+        session.commit()
+    dr = dag.create_dagrun(
+        state=DagRunState.RUNNING,
+        execution_date=execution_date,
+        run_id=run_id,
+        start_date=start_date or execution_date,
+        session=session,
+        conf=conf,  # type: ignore
+    )
+    log.info("created dagrun " + str(dr))
+    return dr

docs/apache-airflow/executor/debug.rst

@@ -15,11 +15,59 @@
     specific language governing permissions and limitations
     under the License.
 
+Testing DAGs with dag.test()
+=============================
+
+To debug DAGs in an IDE, you can set up the ``dag.test`` command in your dag file and run through your DAG in a single
+serialized python process.
+
+This approach can be used with any supported database (including a local SQLite database) and will
+*fail fast* as all tasks run in a single process.
+
+To set up ``dag.test``, add these two lines to the bottom of your dag file:
+
+.. code-block:: python
+
+  if __name__ == "__main__":
+      dag.test()
+
+and that's it! You can add argument such as ``execution_date`` if you want to test argument-specific dagruns, but otherwise
+you can run or debug DAGs as needed.
+
+Comparison with DebugExecutor
+*****************************
+
+The ``dag.test`` command has the following benefits over the :class:`~airflow.executors.debug_executor.DebugExecutor`
+class, which is now deprecated:
+
+1. It does not require running an executor at all. Tasks are run one at a time with no executor or scheduler logs.
+2. It is significantly faster than running code with a DebugExecutor as it does not need to go through a scheduler loop.
+3. It does not perform a backfill.
+
+
+Debugging Airflow DAGs on the command line
+==========================================
+
+With the same two line addition as mentioned in the above section, you can now easily debug a DAG using pdb as well.
+Run ``python -m pdb <path to dag file>.py`` for an interactive debugging experience on the command line.
+
+.. code-block:: bash
+
+  root@ef2c84ad4856:/opt/airflow# python -m pdb airflow/example_dags/example_bash_operator.py
+  > /opt/airflow/airflow/example_dags/example_bash_operator.py(18)<module>()
+  -> """Example DAG demonstrating the usage of the BashOperator."""
+  (Pdb) b 45
+  Breakpoint 1 at /opt/airflow/airflow/example_dags/example_bash_operator.py:45
+  (Pdb) c
+  > /opt/airflow/airflow/example_dags/example_bash_operator.py(45)<module>()
+  -> bash_command='echo 1',
+  (Pdb) run_this_last
+  <Task(EmptyOperator): run_this_last>
 
 .. _executor:DebugExecutor:
 
-Debug Executor
-==================
+Debug Executor (deprecated)
+===========================
 
 The :class:`~airflow.executors.debug_executor.DebugExecutor` is meant as
 a debug tool and can be used from IDE. It is a single process executor that

newsfragments/26400.significant.rst

@@ -0,0 +1,5 @@
+``airflow dags test`` no longer performs a backfill job.
+
+In order to make ``airflow dags test`` more useful as a testing and debugging tool, we no
+longer run a backfill job and instead run a "local task runner". Users can still backfill
+their DAGs using the ``airflow dags backfill`` command.