Repository: apache/airflow

Issue:
DagRun for <FOO> with run_id or execution_date of 'manual__XXXXX' not found
### Apache Airflow version

2.2.2 (latest released)

### What happened

After upgrading from Airflow 2.1.4 to 2.2.2, every DAG gives this error upon execution:

> [2021-12-17, 15:01:12 UTC] {taskinstance.py:1259} INFO - Executing <Task(_PythonDecoratedOperator): print_the_context> on 2021-12-17 15:01:08.943254+00:00
[2021-12-17, 15:01:12 UTC] {standard_task_runner.py:52} INFO - Started process 873 to run task
[2021-12-17, 15:01:12 UTC] {standard_task_runner.py:76} INFO - Running: ['airflow', 'tasks', 'run', 'example_python_operator', 'print_the_context', 'manual__2021-12-17T15:01:08.943254+00:00', '--job-id', '326', '--raw', '--subdir', 'DAGS_FOLDER/test.py', '--cfg-path', '/tmp/tmpej1imvkr', '--error-file', '/tmp/tmpqn9ad7em']
[2021-12-17, 15:01:12 UTC] {standard_task_runner.py:77} INFO - Job 326: Subtask print_the_context
[2021-12-17, 15:01:12 UTC] {standard_task_runner.py:92} ERROR - Failed to execute job 326 for task print_the_context
Traceback (most recent call last):
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/task/task_runner/standard_task_runner.py", line 85, in _start_by_fork
    args.func(args, dag=self.dag)
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/cli/cli_parser.py", line 48, in command
    return func(*args, **kwargs)
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/utils/cli.py", line 92, in wrapper
    return f(*args, **kwargs)
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/cli/commands/task_command.py", line 287, in task_run
    ti = _get_ti(task, args.execution_date_or_run_id)
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/utils/session.py", line 70, in wrapper
    return func(*args, session=session, **kwargs)
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/cli/commands/task_command.py", line 86, in _get_ti
    dag_run = _get_dag_run(task.dag, exec_date_or_run_id, create_if_necssary, session)
  File "/home/ec2-user/venv/lib/python3.7/site-packages/airflow/cli/commands/task_command.py", line 80, in _get_dag_run
    ) from None
airflow.exceptions.DagRunNotFound: DagRun for example_python_operator with run_id or execution_date of 'manual__2021-12-17T15:01:08.943254+00:00' not found

Both tables `airflowdb.task_instance` and `airflowdb.dag_run` have rows with `run_id` equal to "manual__2021-12-17T15:01:08.943254+00:00".

The issue seems to arise in the `_get_dag_run()` function from [airflow/cli/commands/task_command.py](https://github.com/apache/airflow/blob/bb82cc0fbb7a6630eac1155d0c3b445dff13ceb6/airflow/cli/commands/task_command.py#L61-L72):

``` 
    execution_date = None
    with suppress(ParserError, TypeError):
        execution_date = timezone.parse(exec_date_or_run_id)

    if create_if_necessary and not execution_date:
        return DagRun(dag_id=dag.dag_id, run_id=exec_date_or_run_id)
    try:
        return (
            session.query(DagRun)
            .filter(
                DagRun.dag_id == dag.dag_id,
                DagRun.execution_date == execution_date,
            )
            .one()
        )
```
Here, `exec_date_or_run_id == 'manual__2021-12-17T15:01:08.943254+00:00'` and `timezone.parse(exec_date_or_run_id)` fails, meaning `execution_date` stays as `None` and the session query returns no results.

### What you expected to happen

Expect DAGs to run without the above error.

### How to reproduce

Upgraded from 2.1.4 to 2.2.2 and manually ran a few DAGs. The above log is from the [example_python_operator](https://github.com/apache/airflow/blob/main/airflow/example_dags/example_python_operator.py) DAG provided on the Airflow repo.

### Operating System

Amazon Linux 2

### Versions of Apache Airflow Providers

_No response_

### Deployment

Virtualenv installation

### Deployment details

_No response_

### Anything else

Tried `airflow db upgrade` and `airflow db reset`  without any luck. The same issue appears on 2.2.3rc2.

Using MySQL 8.0.23.

### 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:
cc @uranusjr Can you take a look when you have time, please?

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
@@ -16,24 +16,27 @@
 # specific language governing permissions and limitations
 # under the License.
 """Task sub-commands"""
+import datetime
 import importlib
 import json
 import logging
 import os
 import textwrap
-from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress
+from contextlib import contextmanager, redirect_stderr, redirect_stdout
 from typing import List, Optional
 
 from pendulum.parsing.exceptions import ParserError
 from sqlalchemy.orm.exc import NoResultFound
+from sqlalchemy.orm.session import Session
 
 from airflow import settings
 from airflow.cli.simple_table import AirflowConsole
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, DagRunNotFound
+from airflow.exceptions import AirflowException, DagRunNotFound, TaskInstanceNotFound
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.jobs.local_task_job import LocalTaskJob
 from airflow.models import DagPickle, TaskInstance
+from airflow.models.baseoperator import BaseOperator
 from airflow.models.dag import DAG
 from airflow.models.dagrun import DagRun
 from airflow.models.xcom import IN_MEMORY_DAGRUN_ID
@@ -50,46 +53,82 @@
 from airflow.utils.dates import timezone
 from airflow.utils.log.logging_mixin import StreamLogWriter
 from airflow.utils.net import get_hostname
-from airflow.utils.session import create_session, provide_session
-
-
-def _get_dag_run(dag, exec_date_or_run_id, create_if_necessary, session):
+from airflow.utils.session import NEW_SESSION, create_session, provide_session
+
+
+def _get_dag_run(
+    *,
+    dag: DAG,
+    exec_date_or_run_id: str,
+    create_if_necessary: bool,
+    session: Session,
+) -> DagRun:
+    """Try to retrieve a DAG run from a string representing either a run ID or logical date.
+
+    This checks DAG runs like this:
+
+    1. If the input ``exec_date_or_run_id`` matches a DAG run ID, return the run.
+    2. Try to parse the input as a date. If that works, and the resulting
+       date matches a DAG run's logical date, return the run.
+    3. If ``create_if_necessary`` is *False* and the input works for neither of
+       the above, raise ``DagRunNotFound``.
+    4. Try to create a new DAG run. If the input looks like a date, use it as
+       the logical date; otherwise use it as a run ID and set the logical date
+       to the current time.
+    """
     dag_run = dag.get_dagrun(run_id=exec_date_or_run_id, session=session)
     if dag_run:
         return dag_run
 
-    execution_date = None
-    with suppress(ParserError, TypeError):
-        execution_date = timezone.parse(exec_date_or_run_id)
+    try:
+        execution_date: Optional[datetime.datetime] = timezone.parse(exec_date_or_run_id)
+    except (ParserError, TypeError):
+        execution_date = None
 
-    if create_if_necessary and not execution_date:
-        return DagRun(dag_id=dag.dag_id, run_id=exec_date_or_run_id)
     try:
         return (
             session.query(DagRun)
-            .filter(
-                DagRun.dag_id == dag.dag_id,
-                DagRun.execution_date == execution_date,
-            )
+            .filter(DagRun.dag_id == dag.dag_id, DagRun.execution_date == execution_date)
             .one()
         )
     except NoResultFound:
-        if create_if_necessary:
-            return DagRun(dag.dag_id, run_id=IN_MEMORY_DAGRUN_ID, execution_date=execution_date)
-        raise DagRunNotFound(
-            f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
-        ) from None
+        if not create_if_necessary:
+            raise DagRunNotFound(
+                f"DagRun for {dag.dag_id} with run_id or execution_date of {exec_date_or_run_id!r} not found"
+            ) from None
+
+    if execution_date is not None:
+        return DagRun(dag.dag_id, run_id=IN_MEMORY_DAGRUN_ID, execution_date=execution_date)
+    return DagRun(dag.dag_id, run_id=exec_date_or_run_id, execution_date=timezone.utcnow())
 
 
 @provide_session
-def _get_ti(task, exec_date_or_run_id, create_if_necessary=False, session=None):
+def _get_ti(
+    task: BaseOperator,
+    exec_date_or_run_id: str,
+    *,
+    create_if_necessary: bool = False,
+    session: Session = NEW_SESSION,
+) -> TaskInstance:
     """Get the task instance through DagRun.run_id, if that fails, get the TI the old way"""
-    dag_run = _get_dag_run(task.dag, exec_date_or_run_id, create_if_necessary, session)
+    dag_run = _get_dag_run(
+        dag=task.dag,
+        exec_date_or_run_id=exec_date_or_run_id,
+        create_if_necessary=create_if_necessary,
+        session=session,
+    )
 
-    ti = dag_run.get_task_instance(task.task_id)
-    if not ti and create_if_necessary:
+    ti_or_none = dag_run.get_task_instance(task.task_id)
+    if ti_or_none is None:
+        if not create_if_necessary:
+            raise TaskInstanceNotFound(
+                f"TaskInstance for {task.dag.dag_id}, {task.task_id} with "
+                f"run_id or execution_date of {exec_date_or_run_id!r} not found"
+            )
         ti = TaskInstance(task, run_id=dag_run.run_id)
         ti.dag_run = dag_run
+    else:
+        ti = ti_or_none
     ti.refresh_from_task(task)
     return ti
 


Test patch: