Repository: apache/airflow

Issue:
current_state method of TaskInstance fails for mapped task instance
### Apache Airflow version

2.4.3

### What happened

`current_state` method on TaskInstance doesn't filter by `map_index` so calling this method on mapped task instance fails.

https://github.com/apache/airflow/blob/fb7c6afc8cb7f93909bd2e654ea185eb6abcc1ea/airflow/models/taskinstance.py#L708-L726

### What you think should happen instead

map_index should also be filtered in the query to return single TaskInstance object.

### How to reproduce

```python
with create_session() as session:
    print(session.query(TaskInstance).filter(TaskInstance.dag_id == "divide_by_zero", 
                                             TaskInstance.map_index == 1, 
                                             TaskInstance.run_id == 'scheduled__2022-11-22T00:00:00+00:00')
                                     .scalar().current_state())

 
---------------------------------------------------------------------------
MultipleResultsFound                      Traceback (most recent call last)
Input In [7], in <cell line: 1>()
      1 with create_session() as session:
----> 2     print(session.query(TaskInstance).filter(TaskInstance.dag_id == "divide_by_zero", TaskInstance.map_index == 1, TaskInstance.run_id == 'scheduled__2022-11-22T00:00:00+00:00').scalar().current_state())

File ~/stuff/python/airflow/airflow/utils/session.py:75, in provide_session.<locals>.wrapper(*args, **kwargs)
     73 else:
     74     with create_session() as session:
---> 75         return func(*args, session=session, **kwargs)

File ~/stuff/python/airflow/airflow/models/taskinstance.py:725, in TaskInstance.current_state(self, session)
    708 @provide_session
    709 def current_state(self, session: Session = NEW_SESSION) -> str:
    710     """
    711     Get the very latest state from the database, if a session is passed,
    712     we use and looking up the state becomes part of the session, otherwise
   (...)
    715     :param session: SQLAlchemy ORM Session
    716     """
    717     return (
    718         session.query(TaskInstance.state)
    719         .filter(
    720             TaskInstance.dag_id == self.dag_id,
    721             TaskInstance.task_id == self.task_id,
    722             TaskInstance.run_id == self.run_id,
    723
    724         )
--> 725         .scalar()
    726     )

File ~/stuff/python/airflow/.env/lib/python3.10/site-packages/sqlalchemy/orm/query.py:2803, in Query.scalar(self)
   2801 # TODO: not sure why we can't use result.scalar() here
   2802 try:
-> 2803     ret = self.one()
   2804     if not isinstance(ret, collections_abc.Sequence):
   2805         return ret

File ~/stuff/python/airflow/.env/lib/python3.10/site-packages/sqlalchemy/orm/query.py:2780, in Query.one(self)
   2762 def one(self):
   2763     """Return exactly one result or raise an exception.
   2764 
   2765     Raises ``sqlalchemy.orm.exc.NoResultFound`` if the query selects
   (...)
   2778 
   2779     """
-> 2780     return self._iter().one()

File ~/stuff/python/airflow/.env/lib/python3.10/site-packages/sqlalchemy/engine/result.py:1162, in Result.one(self)
   1134 def one(self):
   1135     # type: () -> Row
   1136     """Return exactly one row or raise an exception.
   1137 
   1138     Raises :class:`.NoResultFound` if the result returns no
   (...)
   1160 
   1161     """
-> 1162     return self._only_one_row(True, True, False)

File ~/stuff/python/airflow/.env/lib/python3.10/site-packages/sqlalchemy/engine/result.py:620, in ResultInternal._only_one_row(self, raise_for_second_row, raise_for_none, scalar)
    618     if next_row is not _NO_ROW:
    619         self._soft_close(hard=True)
--> 620         raise exc.MultipleResultsFound(
    621             "Multiple rows were found when exactly one was required"
    622             if raise_for_none
    623             else "Multiple rows were found when one or none "
    624             "was required"
    625         )
    626 else:
    627     next_row = _NO_ROW

MultipleResultsFound: Multiple rows were found when exactly one was required

```

### Operating System

Ubuntu

### Versions of Apache Airflow Providers

_No response_

### Deployment

Virtualenv installation

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [X] 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:

Resolving patch:
diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -736,17 +736,13 @@ def current_state(self, session: Session = NEW_SESSION) -> str:
         we use and looking up the state becomes part of the session, otherwise
         a new session is used.
 
+        sqlalchemy.inspect is used here to get the primary keys ensuring that if they change
+        it will not regress
+
         :param session: SQLAlchemy ORM Session
         """
-        return (
-            session.query(TaskInstance.state)
-            .filter(
-                TaskInstance.dag_id == self.dag_id,
-                TaskInstance.task_id == self.task_id,
-                TaskInstance.run_id == self.run_id,
-            )
-            .scalar()
-        )
+        filters = (col == getattr(self, col.name) for col in inspect(TaskInstance).primary_key)
+        return session.query(TaskInstance.state).filter(*filters).scalar()
 
     @provide_session
     def error(self, session: Session = NEW_SESSION) -> None:


Test patch:
diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py
--- a/tests/models/test_taskinstance.py
+++ b/tests/models/test_taskinstance.py
@@ -1863,6 +1863,29 @@ def test_outlet_datasets_failed(self, create_task_instance):
         # check that no dataset events were generated
         assert session.query(DatasetEvent).count() == 0
 
+    def test_mapped_current_state(self, dag_maker):
+        with dag_maker(dag_id="test_mapped_current_state") as _:
+            from airflow.decorators import task
+
+            @task()
+            def raise_an_exception(placeholder: int):
+                if placeholder == 0:
+                    raise AirflowFailException("failing task")
+                else:
+                    pass
+
+            _ = raise_an_exception.expand(placeholder=[0, 1])
+
+        tis = dag_maker.create_dagrun(execution_date=timezone.utcnow()).task_instances
+        for task_instance in tis:
+            if task_instance.map_index == 0:
+                with pytest.raises(AirflowFailException):
+                    task_instance.run()
+                assert task_instance.current_state() == TaskInstanceState.FAILED
+            else:
+                task_instance.run()
+                assert task_instance.current_state() == TaskInstanceState.SUCCESS
+
     def test_outlet_datasets_skipped(self, create_task_instance):
         """
         Verify that when we have an outlet dataset on a task, and the task