From 2b30c53834072dc251ad6a24d38c643efad57790 Mon Sep 17 00:00:00 2001
From: Kaxil Naik <kaxilnaik@gmail.com>
Date: Mon, 25 Jan 2021 23:01:13 +0000
Subject: [PATCH] Improve the error when DAG does not exist when running dag
 pause command

When running `airflow dags unpause` with a DAG that does not exist, it
currently shows this error

```
root@6f086ba87198:/opt/airflow# airflow dags unpause example_bash_operatoredd
Traceback (most recent call last):
  File "/usr/local/bin/airflow", line 33, in <module>
    sys.exit(load_entry_point('apache-airflow', 'console_scripts', 'airflow')())
  File "/opt/airflow/airflow/__main__.py", line 40, in main
    args.func(args)
  File "/opt/airflow/airflow/cli/cli_parser.py", line 48, in command
    return func(*args, **kwargs)
  File "/opt/airflow/airflow/utils/cli.py", line 92, in wrapper
    return f(*args, **kwargs)
  File "/opt/airflow/airflow/cli/commands/dag_command.py", line 160, in dag_unpause
    set_is_paused(False, args)
  File "/opt/airflow/airflow/cli/commands/dag_command.py", line 170, in set_is_paused
    dag.set_is_paused(is_paused=is_paused)
AttributeError: 'NoneType' object has no attribute 'set_is_paused'
```

This commit changes the error to show a helpful error:

```
root@6f086ba87198:/opt/airflow# airflow dags unpause example_bash_operatoredd
DAG: example_bash_operatoredd does not exit in 'dag' table
```
---
 airflow/cli/commands/dag_command.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py
index 48b04df77e8a1..b1fa80170a13d 100644
--- a/airflow/cli/commands/dag_command.py
+++ b/airflow/cli/commands/dag_command.py
@@ -162,9 +162,12 @@ def dag_unpause(args):
 
 def set_is_paused(is_paused, args):
     """Sets is_paused for DAG by a given dag_id"""
-    DagModel.get_dagmodel(args.dag_id).set_is_paused(
-        is_paused=is_paused,
-    )
+    dag = DagModel.get_dagmodel(args.dag_id)
+
+    if not dag:
+        raise SystemExit(f"DAG: {args.dag_id} does not exit in 'dag' table")
+
+    dag.set_is_paused(is_paused=is_paused)
 
     print(f"Dag: {args.dag_id}, paused: {is_paused}")