Repository: apache/airflow

Issue:
All task logging goes to the log for try_number 1
**Apache Airflow version**: 2.0.0a1
**What happened**:

When a task fails on the first try, the log output for additional tries go to the log for the first attempt.

**What you expected to happen**:

The logs should go to the correct log file. For the default configuration, the log filename template is `log_filename_template = {{ ti.dag_id }}/{{ ti.task_id }}/{{ ts }}/{{ try_number }}.log`, so additional numbered `.log` files should be created.

**How to reproduce it**:

Create a test dag:

```
from datetime import timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago


with DAG(
    dag_id="trynumber_demo",
    default_args={"start_date": days_ago(2), "retries": 1, "retry_delay": timedelta(0)},
    schedule_interval=None,
) as dag:

    def demo_task(ti=None):
        print("Running demo_task, try_number =", ti.try_number)
        if ti.try_number <= 1:
            raise ValueError("Shan't")

    task = PythonOperator(task_id="demo_task", python_callable=demo_task)
```

and trigger this dag:

```
$ airflow dags trigger trynumber_demo
```

then observe that `triggernumber_demo/demo_task/<execution_date>/` only contains 1.log, which contains the full output for 2 runs:

```
[...]
--------------------------------------------------------------------------------
[2020-10-21 13:29:07,958] {taskinstance.py:1020} INFO - Starting attempt 1 of 2
[2020-10-21 13:29:07,959] {taskinstance.py:1021} INFO -
--------------------------------------------------------------------------------
[...]
[2020-10-21 13:29:08,163] {logging_mixin.py:110} INFO - Running demo_task, try_number = 1
[2020-10-21 13:29:08,164] {taskinstance.py:1348} ERROR - Shan't
Traceback (most recent call last):
[...]
ValueError: Shan't
[2020-10-21 13:29:08,168] {taskinstance.py:1392} INFO - Marking task as UP_FOR_RETRY. dag_id=trynumber_demo, task_id=demo_task, execution_date=20201021T122907, start_date=20201021T122907, end_date=20201021T122908
[...]
[2020-10-21 13:29:09,121] {taskinstance.py:1019} INFO -
--------------------------------------------------------------------------------
[2020-10-21 13:29:09,121] {taskinstance.py:1020} INFO - Starting attempt 2 of 2
[2020-10-21 13:29:09,121] {taskinstance.py:1021} INFO -
--------------------------------------------------------------------------------
[...]
[2020-10-21 13:29:09,333] {logging_mixin.py:110} INFO - Running demo_task, try_number = 2
[2020-10-21 13:29:09,334] {python.py:141} INFO - Done. Returned value was: None
[2020-10-21 13:29:09,355] {taskinstance.py:1143} INFO - Marking task as SUCCESS.dag_id=trynumber_demo, task_id=demo_task, execution_date=20201021T122907, start_date=20201021T122909, end_date=20201021T122909
[2020-10-21 13:29:09,404] {local_task_job.py:117} INFO - Task exited with return code 0
```

The `TaskInstance()` created for the run needs to first be refreshed from the database, before setting the logging context.


Discussion:
I can confirm that changing

https://github.com/apache/airflow/blob/172820db4d2009dd26fa8aef4a864fb8a3d7e78d/airflow/cli/commands/task_command.py#L172-L174

to

```
    task = dag.get_task(task_id=args.task_id)
    ti = TaskInstance(task, args.execution_date)
    ti.refresh_from_db()
    ti.init_run_context(raw=args.raw)
```

fixes this issue, but I don't know if the `refresh_..` call should be integrated into `init_run_context()` perhaps.
I found #7370, which removed that call, specifically see [these review comments](https://github.com/apache/airflow/pull/7370#discussion_r375370760):

> Does init_run_context look at any attributes of the TI? If so they won't be populated until this call.

> No, it only setups logging context :)

Unfortunately, the logging context looks at the attributes of the TI!
@mjpieters would you like to open a PR?
Could be same as https://github.com/apache/airflow/issues/11561 (though more detail/cause here)

Resolving patch:
diff --git a/airflow/cli/commands/task_command.py b/airflow/cli/commands/task_command.py
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -171,6 +171,7 @@ def task_run(args, dag=None):
 
     task = dag.get_task(task_id=args.task_id)
     ti = TaskInstance(task, args.execution_date)
+    ti.refresh_from_db()
     ti.init_run_context(raw=args.raw)
 
     hostname = get_hostname()


Test patch:
diff --git a/tests/cli/commands/test_task_command.py b/tests/cli/commands/test_task_command.py
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -31,11 +31,12 @@
 from airflow.cli.commands import task_command
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import DagBag, TaskInstance
-from airflow.settings import Session
+from airflow.models import DagBag, DagRun, TaskInstance
 from airflow.utils import timezone
 from airflow.utils.cli import get_dag
+from airflow.utils.session import create_session
 from airflow.utils.state import State
+from airflow.utils.types import DagRunType
 from tests.test_utils.config import conf_vars
 from tests.test_utils.db import clear_db_pools, clear_db_runs
 
@@ -46,11 +47,11 @@
 
 
 def reset(dag_id):
-    session = Session()
-    tis = session.query(TaskInstance).filter_by(dag_id=dag_id)
-    tis.delete()
-    session.commit()
-    session.close()
+    with create_session() as session:
+        tis = session.query(TaskInstance).filter_by(dag_id=dag_id)
+        tis.delete()
+        runs = session.query(DagRun).filter_by(dag_id=dag_id)
+        runs.delete()
 
 
 class TestCliTasks(unittest.TestCase):
@@ -256,8 +257,10 @@ class TestLogsfromTaskRunCommand(unittest.TestCase):
     def setUp(self) -> None:
         self.dag_id = "test_logging_dag"
         self.task_id = "test_task"
+        self.dag_path = os.path.join(ROOT_FOLDER, "dags", "test_logging_in_dag.py")
         reset(self.dag_id)
-        self.execution_date_str = timezone.make_aware(datetime(2017, 1, 1)).isoformat()
+        self.execution_date = timezone.make_aware(datetime(2017, 1, 1))
+        self.execution_date_str = self.execution_date.isoformat()
         self.log_dir = conf.get('logging', 'base_log_folder')
         self.log_filename = f"{self.dag_id}/{self.task_id}/{self.execution_date_str}/1.log"
         self.ti_log_file_path = os.path.join(self.log_dir, self.log_filename)
@@ -295,7 +298,7 @@ def test_logging_with_run_task(self):
         #  We are not using self.assertLogs as we want to verify what actually is stored in the Log file
         # as that is what gets displayed
 
-        with conf_vars({('core', 'dags_folder'): os.path.join(ROOT_FOLDER, f"tests/dags/{self.dag_id}")}):
+        with conf_vars({('core', 'dags_folder'): self.dag_path}):
             task_command.task_run(self.parser.parse_args([
                 'tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]))
 
@@ -322,7 +325,7 @@ def test_logging_with_run_task(self):
     def test_logging_with_run_task_subprocess(self):
         # We are not using self.assertLogs as we want to verify what actually is stored in the Log file
         # as that is what gets displayed
-        with conf_vars({('core', 'dags_folder'): os.path.join(ROOT_FOLDER, f"tests/dags/{self.dag_id}")}):
+        with conf_vars({('core', 'dags_folder'): self.dag_path}):
             task_command.task_run(self.parser.parse_args([
                 'tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]))
 
@@ -343,6 +346,43 @@ def test_logging_with_run_task_subprocess(self):
         self.assertIn(f"INFO - Marking task as SUCCESS.dag_id={self.dag_id}, "
                       f"task_id={self.task_id}, execution_date=20170101T000000", logs)
 
+    def test_log_file_template_with_run_task(self):
+        """Verify that the taskinstance has the right context for log_filename_template"""
+
+        with mock.patch.object(task_command, "_run_task_by_selected_method"):
+            with conf_vars({('core', 'dags_folder'): self.dag_path}):
+                # increment the try_number of the task to be run
+                dag = DagBag().get_dag(self.dag_id)
+                task = dag.get_task(self.task_id)
+                with create_session() as session:
+                    dag.create_dagrun(
+                        execution_date=self.execution_date,
+                        start_date=timezone.utcnow(),
+                        state=State.RUNNING,
+                        run_type=DagRunType.MANUAL,
+                        session=session,
+                    )
+                    ti = TaskInstance(task, self.execution_date)
+                    ti.refresh_from_db(session=session, lock_for_update=True)
+                    ti.try_number = 1  # not running, so starts at 0
+                    session.merge(ti)
+
+                log_file_path = os.path.join(os.path.dirname(self.ti_log_file_path), "2.log")
+
+                try:
+                    task_command.task_run(
+                        self.parser.parse_args(
+                            ['tasks', 'run', self.dag_id, self.task_id, '--local', self.execution_date_str]
+                        )
+                    )
+
+                    assert os.path.exists(log_file_path)
+                finally:
+                    try:
+                        os.remove(log_file_path)
+                    except OSError:
+                        pass
+
 
 class TestCliTaskBackfill(unittest.TestCase):
     @classmethod