Move fallible ti.task.dag assignment back inside try/except block (#2…

…4533) (#24592)

* Move fallible ti.task.dag assignment back inside try/except block

It looks like ti.task.dag was originally protected inside try/except,
but was moved out at commit 7be87d

* Remove unneeded variable annotation

Co-authored-by: EJ Kreinar <ej@he360.com>
Co-authored-by: Tzu-ping Chung <uranusjr@gmail.com>

airflow/providers/elasticsearch/log/es_task_handler.py

@@ -125,12 +125,13 @@ def _render_log_id(self, ti: TaskInstance, try_number: int) -> str:
             else:
                 log_id_template = self.log_id_template
 
-        dag = ti.task.dag
-        assert dag is not None  # For Mypy.
         try:
-            data_interval: Tuple[datetime, datetime] = dag.get_run_data_interval(dag_run)
+            dag = ti.task.dag
         except AttributeError:  # ti.task is not always set.
             data_interval = (dag_run.data_interval_start, dag_run.data_interval_end)
+        else:
+            assert dag is not None  # For Mypy.
+            data_interval = dag.get_run_data_interval(dag_run)
 
         if self.json_format:
             data_interval_start = self._clean_date(data_interval[0])

airflow/utils/log/file_task_handler.py

@@ -19,9 +19,8 @@
 import logging
 import os
 import warnings
-from datetime import datetime
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Optional
 
 from airflow.configuration import AirflowConfigException, conf
 from airflow.utils.context import Context
@@ -93,12 +92,13 @@ def _render_filename(self, ti: "TaskInstance", try_number: int) -> str:
             context["try_number"] = try_number
             return render_template_to_string(jinja_tpl, context)
         elif str_tpl:
-            dag = ti.task.dag
-            assert dag is not None  # For Mypy.
             try:
-                data_interval: Tuple[datetime, datetime] = dag.get_run_data_interval(dag_run)
+                dag = ti.task.dag
             except AttributeError:  # ti.task is not always set.
                 data_interval = (dag_run.data_interval_start, dag_run.data_interval_end)
+            else:
+                assert dag is not None  # For Mypy.
+                data_interval = dag.get_run_data_interval(dag_run)
             if data_interval[0]:
                 data_interval_start = data_interval[0].isoformat()
             else: