Fallback to parse dag_file when no dag in the db (apache#23738)

NayaraWakewski · May 18, 2022 · f721529 · f721529
1 parent 5d743e7
commit f721529
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 1 deletion.
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
@@ -358,7 +358,11 @@ def task_run(args, dag=None):
         dag = get_dag_by_pickle(args.pickle)
     elif not dag:
         if args.local:
-            dag = get_dag_by_deserialization(args.dag_id)
+            try:
+                dag = get_dag_by_deserialization(args.dag_id)
+            except AirflowException:
+                print(f'DAG {args.dag_id} does not exist in the database, trying to parse the dag_file')
+                dag = get_dag(args.subdir, args.dag_id)
         else:
             dag = get_dag(args.subdir, args.dag_id)
     else:

diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
@@ -152,6 +152,38 @@ def test_run_get_serialized_dag(self, mock_local_job, mock_get_dag_by_deserializ
         )
         mock_get_dag_by_deserialization.assert_called_once_with(self.dag_id)
 
+    @mock.patch("airflow.cli.commands.task_command.get_dag_by_deserialization")
+    @mock.patch("airflow.cli.commands.task_command.LocalTaskJob")
+    def test_run_get_serialized_dag_fallback(self, mock_local_job, mock_get_dag_by_deserialization):
+        """
+        Fallback to parse dag_file when serialized dag does not exist in the db
+        """
+        task_id = self.dag.task_ids[0]
+        args = [
+            'tasks',
+            'run',
+            '--ignore-all-dependencies',
+            '--local',
+            self.dag_id,
+            task_id,
+            self.run_id,
+        ]
+        mock_get_dag_by_deserialization.side_effect = mock.Mock(side_effect=AirflowException('Not found'))
+
+        task_command.task_run(self.parser.parse_args(args))
+        mock_local_job.assert_called_once_with(
+            task_instance=mock.ANY,
+            mark_success=False,
+            ignore_all_deps=True,
+            ignore_depends_on_past=False,
+            ignore_task_deps=False,
+            ignore_ti_state=False,
+            pickle_id=None,
+            pool=None,
+            external_executor_id=None,
+        )
+        mock_get_dag_by_deserialization.assert_called_once_with(self.dag_id)
+
     @mock.patch("airflow.cli.commands.task_command.LocalTaskJob")
     def test_run_with_existing_dag_run_id(self, mock_local_job):
         """