Repository: apache/airflow

Issue:
`airflow dags status` fails if parse time is near `dagbag_import_timeout`
### Apache Airflow version

2.2.3 (latest released)

### What happened

I had just kicked off a DAG and I was periodically running `airflow dags status ...` to see if it was done yet.  At first it seemed to work, but later it failed with this error:

```
$ airflow dags state load_13 '2022-02-09T05:25:28+00:00'

    [2022-02-09 05:26:56,493] {dagbag.py:500} INFO - Filling up the DagBag from /usr/local/airflow/dags
    queued

$ airflow dags state load_13 '2022-02-09T05:25:28+00:00'

    [2022-02-09 05:27:29,096] {dagbag.py:500} INFO - Filling up the DagBag from /usr/local/airflow/dags
    [2022-02-09 05:27:59,084] {timeout.py:36} ERROR - Process timed out, PID: 759
    [2022-02-09 05:27:59,088] {dagbag.py:334} ERROR - Failed to import: /usr/local/airflow/dags/many_tasks.py
    Traceback (most recent call last):
      File "/usr/local/lib/python3.9/site-packages/airflow/models/dagbag.py", line 331, in _load_modules_from_file
        loader.exec_module(new_module)
      File "<frozen importlib._bootstrap_external>", line 850, in exec_module
      File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
      File "/usr/local/airflow/dags/many_tasks.py", line 61, in <module>
        globals()["dag_{:02d}".format(i)] = parameterized_load(i, step)
      File "/usr/local/airflow/dags/many_tasks.py", line 50, in parameterized_load
        return load()
      File "/usr/local/lib/python3.9/site-packages/airflow/models/dag.py", line 2984, in factory
        f(**f_kwargs)
      File "/usr/local/airflow/dags/many_tasks.py", line 48, in load
        [worker_factory(i) for i in range(1, size**2 + 1)]
      File "/usr/local/airflow/dags/many_tasks.py", line 48, in <listcomp>
        [worker_factory(i) for i in range(1, size**2 + 1)]
      File "/usr/local/airflow/dags/many_tasks.py", line 37, in worker_factory
        return worker(num)
      File "/usr/local/lib/python3.9/site-packages/airflow/decorators/base.py", line 219, in factory
        op = decorated_operator_class(
      File "/usr/local/lib/python3.9/site-packages/airflow/models/baseoperator.py", line 188, in apply_defaults
        result = func(self, *args, **kwargs)
      File "/usr/local/lib/python3.9/site-packages/airflow/decorators/python.py", line 59, in __init__
        super().__init__(kwargs_to_upstream=kwargs_to_upstream, **kwargs)
      File "/usr/local/lib/python3.9/site-packages/airflow/models/baseoperator.py", line 152, in apply_defaults
        dag_params = copy.deepcopy(dag.params) or {}
      File "/usr/local/lib/python3.9/copy.py", line 172, in deepcopy
        y = _reconstruct(x, memo, *rv)
      File "/usr/local/lib/python3.9/copy.py", line 264, in _reconstruct
        y = func(*args)
      File "/usr/local/lib/python3.9/copy.py", line 263, in <genexpr>
        args = (deepcopy(arg, memo) for arg in args)
      File "/usr/local/lib/python3.9/site-packages/airflow/utils/timeout.py", line 37, in handle_timeout
        raise AirflowTaskTimeout(self.error_message)
    airflow.exceptions.AirflowTaskTimeout: DagBag import timeout for /usr/local/airflow/dags/many_tasks.py after 30.0s.
    Please take a look at these docs to improve your DAG import time:
    * http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/latest/best-practices.html#top-level-python-code
    * http://apache-airflow-docs.s3-website.eu-central-1.amazonaws.com/docs/apache-airflow/latest/best-practices.html#reducing-dag-complexity, PID: 759
    Traceback (most recent call last):
      File "/usr/local/bin/airflow", line 8, in <module>
        sys.exit(main())
      File "/usr/local/lib/python3.9/site-packages/airflow/__main__.py", line 48, in main
        args.func(args)
      File "/usr/local/lib/python3.9/site-packages/airflow/cli/cli_parser.py", line 48, in command
        return func(*args, **kwargs)
      File "/usr/local/lib/python3.9/site-packages/airflow/utils/cli.py", line 92, in wrapper
        return f(*args, **kwargs)
      File "/usr/local/lib/python3.9/site-packages/airflow/cli/commands/dag_command.py", line 241, in dag_state
        dag = get_dag(args.subdir, args.dag_id)
      File "/usr/local/lib/python3.9/site-packages/airflow/utils/cli.py", line 192, in get_dag
        raise AirflowException(
    airflow.exceptions.AirflowException: Dag 'load_13' could not be found; either it does not exist or it failed to parse.
```

### What you expected to happen

If we were able to parse the DAG in the first place, I expect that downstream actions (like querying for status) would not fail due to a dag parsing timeout.

Also, is parsing the dag necessary for this action?

### How to reproduce

1. start with the dag shown here: https://gist.github.com/MatrixManAtYrService/842266aac42390aadee75fe014cd372e
2. increase "scale" until `airflow dags list` stop showing the load dags
3. decrease by one and check that they start showing back up
4. trigger a dag run
5. check its status (periodically), eventually the status check will fail

I initially discovered this using the `CeleryExecutor` and a much messier dag, but once I understood what I was looking for I was able to recreate it using the dag linked above and `astro dev start`

### Operating System

docker/debian

### Versions of Apache Airflow Providers

n/a

### Deployment

Astronomer

### Deployment details

```
FROM quay.io/astronomer/ap-airflow:2.2.3-onbuild
```

### Anything else

When I was running this via the CeleryExecutor (deployed via helm on a single-node k8s cluster), I noticed similar dag-parsing timeouts showing up in the worker logs.  I failed to capture them because I didn't yet know what I was looking for, but if they would be helpful I can recreate that scenario and post them here.

----

I tried to work around this error by doubling the following configs:
- [dagbag_import_timeout](https://airflow.apache.org/docs/apache-airflow/stable/configurations-ref.html#dagbag-import-timeout)
- [dag_file_processor_timeout](https://airflow.apache.org/docs/apache-airflow/stable/configurations-ref.html#dag-file-processor-timeout)

This "worked", as in the status started showing up without error, but it seemed like making the dag __longer__ had also made it  __slower__.  As if whatever re-parsing steps were occurring along the way were also slowing it down.  It used to take 1h to complete, but when I checked on it after 1h it was only 30% complete (the new tasks hadn't even started yet).

Should I expect that splitting my large dag into smaller dags will fix this?  Or is the overall parsing workload going to eat into my runtime regardless of how it is sliced?

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
Yes. Parsing is necessary to get the structure of DAGs. The command line of airflow is not intended as "external" query mechanism, it is a "helper" tool that should be executed on Airflow scheduler and it is supposed to read the DB and parse DAGs. 

If you want to query dag state from outside use the REST API: https://airflow.apache.org/docs/apache-airflow/stable/stable-rest-api-ref.html
Parsing the file is not necessary -- it is possible for this command to use the DB -- after all, all it is  looking at is the status of a DagRun.

Resolving patch:
diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py
--- a/airflow/cli/commands/dag_command.py
+++ b/airflow/cli/commands/dag_command.py
@@ -39,15 +39,9 @@
 from airflow.models.dag import DAG
 from airflow.models.serialized_dag import SerializedDagModel
 from airflow.utils import cli as cli_utils
-from airflow.utils.cli import (
-    get_dag,
-    get_dag_by_file_location,
-    process_subdir,
-    sigint_handler,
-    suppress_logs_and_warning,
-)
+from airflow.utils.cli import get_dag, process_subdir, sigint_handler, suppress_logs_and_warning
 from airflow.utils.dot_renderer import render_dag, render_dag_dependencies
-from airflow.utils.session import create_session, provide_session
+from airflow.utils.session import NEW_SESSION, create_session, provide_session
 from airflow.utils.state import DagRunState
 
 
@@ -249,7 +243,8 @@ def _save_dot_to_file(dot: Dot, filename: str):
 
 
 @cli_utils.action_cli
-def dag_state(args):
+@provide_session
+def dag_state(args, session=NEW_SESSION):
     """
     Returns the state (and conf if exists) of a DagRun at the command line.
     >>> airflow dags state tutorial 2015-01-01T00:00:00.000000
@@ -257,15 +252,16 @@ def dag_state(args):
     >>> airflow dags state a_dag_with_conf_passed 2015-01-01T00:00:00.000000
     failed, {"name": "bob", "age": "42"}
     """
-    if args.subdir:
-        dag = get_dag(args.subdir, args.dag_id)
-    else:
-        dag = get_dag_by_file_location(args.dag_id)
-    dr = DagRun.find(dag.dag_id, execution_date=args.execution_date)
-    out = dr[0].state if dr else None
+
+    dag = DagModel.get_dagmodel(args.dag_id, session=session)
+
+    if not dag:
+        raise SystemExit(f"DAG: {args.dag_id} does not exist in 'dag' table")
+    dr = session.query(DagRun).filter_by(dag_id=args.dag_id, execution_date=args.execution_date).one_or_none()
+    out = dr.state if dr else None
     conf_out = ''
-    if out and dr[0].conf:
-        conf_out = ', ' + json.dumps(dr[0].conf)
+    if out and dr.conf:
+        conf_out = ', ' + json.dumps(dr.conf)
     print(str(out) + conf_out)
 
 
@@ -351,32 +347,27 @@ def dag_report(args):
 
 @cli_utils.action_cli
 @suppress_logs_and_warning
-def dag_list_jobs(args, dag=None):
+@provide_session
+def dag_list_jobs(args, dag=None, session=NEW_SESSION):
     """Lists latest n jobs"""
     queries = []
     if dag:
         args.dag_id = dag.dag_id
     if args.dag_id:
-        dagbag = DagBag()
+        dag = DagModel.get_dagmodel(args.dag_id, session=session)
 
-        if args.dag_id not in dagbag.dags:
-            error_message = f"Dag id {args.dag_id} not found"
-            raise AirflowException(error_message)
+        if not dag:
+            raise SystemExit(f"DAG: {args.dag_id} does not exist in 'dag' table")
         queries.append(BaseJob.dag_id == args.dag_id)
 
     if args.state:
         queries.append(BaseJob.state == args.state)
 
     fields = ['dag_id', 'state', 'job_type', 'start_date', 'end_date']
-    with create_session() as session:
-        all_jobs = (
-            session.query(BaseJob)
-            .filter(*queries)
-            .order_by(BaseJob.start_date.desc())
-            .limit(args.limit)
-            .all()
-        )
-        all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs]
+    all_jobs = (
+        session.query(BaseJob).filter(*queries).order_by(BaseJob.start_date.desc()).limit(args.limit).all()
+    )
+    all_jobs = [{f: str(job.__getattribute__(f)) for f in fields} for job in all_jobs]
 
     AirflowConsole().print_as(
         data=all_jobs,
@@ -386,16 +377,16 @@ def dag_list_jobs(args, dag=None):
 
 @cli_utils.action_cli
 @suppress_logs_and_warning
-def dag_list_dag_runs(args, dag=None):
+@provide_session
+def dag_list_dag_runs(args, dag=None, session=NEW_SESSION):
     """Lists dag runs for a given DAG"""
     if dag:
         args.dag_id = dag.dag_id
+    else:
+        dag = DagModel.get_dagmodel(args.dag_id, session=session)
 
-    dagbag = DagBag()
-
-    if args.dag_id is not None and args.dag_id not in dagbag.dags:
-        error_message = f"Dag id {args.dag_id} not found"
-        raise AirflowException(error_message)
+        if not dag:
+            raise SystemExit(f"DAG: {args.dag_id} does not exist in 'dag' table")
 
     state = args.state.lower() if args.state else None
     dag_runs = DagRun.find(
@@ -404,6 +395,7 @@ def dag_list_dag_runs(args, dag=None):
         no_backfills=args.no_backfill,
         execution_start_date=args.start_date,
         execution_end_date=args.end_date,
+        session=session,
     )
 
     dag_runs.sort(key=lambda x: x.execution_date, reverse=True)


Test patch: