fixup! Read current task executions by exact region and try

Author: ashb

airflow-core/src/airflow/api_fastapi/core_api/routes/public/hitl.py

@@ -115,8 +115,7 @@ def _get_task_instance_with_hitl_detail(
         dag_id=dag_id,
         run_id=dag_run_id,
         task_id=task_id,
-        session=session,
-        dag_bag=dag_bag,
+        resolver=TaskCoordinateResolver(dag_bag, session),
         map_index=map_index,
         region_id=region_id,
         region_index=region_index,

airflow-core/src/airflow/api_fastapi/core_api/routes/public/log.py

@@ -211,6 +211,7 @@ def get_external_log_url(
             TaskInstance.try_number == try_number,
         )
         .options(joinedload(TaskInstance.dag_model))
+        .execution_options(include_all_attempts=True)
     )
     ti = session.scalar(query)
 

airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py

@@ -26,7 +26,6 @@
 
 from airflow._shared.state import TaskScope
 from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity
-from airflow.api_fastapi.common.dagbag import DagBagDep
 from airflow.api_fastapi.common.db.common import SessionDep, paginated_select
 from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset
 from airflow.api_fastapi.common.router import AirflowRouter
@@ -38,7 +37,11 @@
 )
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
 from airflow.api_fastapi.core_api.security import requires_access_dag
-from airflow.api_fastapi.core_api.services.public.task_coordinates import TaskScopeDep, resolve_task_scope
+from airflow.api_fastapi.core_api.services.public.task_coordinates import (
+    CoordinateResolverDep,
+    TaskScopeDep,
+    resolve_task_scope,
+)
 from airflow.configuration import conf
 from airflow.models.dynamic_region import SENTINEL_REGION_ID, DynamicRegion
 from airflow.models.task_state_store import TaskStateStoreModel
@@ -97,8 +100,7 @@ def _resolve_clear_scope(
     dag_id: str,
     dag_run_id: str,
     task_id: str,
-    session: SessionDep,
-    dag_bag: DagBagDep,
+    resolver: CoordinateResolverDep,
     map_index: Annotated[int, Query(ge=-1)] = -1,
     region_id: Annotated[UUID | None, Query()] = None,
     region_index: Annotated[int | None, Query(ge=-1)] = None,
@@ -108,8 +110,7 @@ def _resolve_clear_scope(
         dag_id=dag_id,
         run_id=dag_run_id,
         task_id=task_id,
-        session=session,
-        dag_bag=dag_bag,
+        resolver=resolver,
         map_index=map_index,
         region_id=region_id,
         region_index=region_index,

airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py

@@ -62,7 +62,7 @@
 from airflow.exceptions import TaskNotFound
 from airflow.models import DagRun as DR
 from airflow.models.dag import DagModel
-from airflow.models.task_coordinates import public_map_index_expression
+from airflow.models.task_coordinates import TaskCoordinateResolver, public_map_index_expression
 from airflow.models.taskinstance import TaskInstance
 from airflow.models.xcom import XComModel, build_xcom_read_query, select_producers, xcom_entity
 
@@ -344,8 +344,7 @@ def create_xcom_entry(
         dag_id=dag_id,
         run_id=dag_run_id,
         task_id=task_id,
-        session=session,
-        dag_bag=dag_bag,
+        resolver=TaskCoordinateResolver(dag_bag, session),
         map_index=request_body.map_index,
         region_id=request_body.region_id,
         region_index=request_body.region_index,
@@ -430,8 +429,7 @@ def update_xcom_entry(
         dag_id=dag_id,
         run_id=dag_run_id,
         task_id=task_id,
-        session=session,
-        dag_bag=dag_bag,
+        resolver=TaskCoordinateResolver(dag_bag, session),
         map_index=patch_body.map_index,
         region_id=patch_body.region_id,
         region_index=patch_body.region_index,

airflow-core/src/airflow/api_fastapi/core_api/services/public/task_coordinates.py

@@ -30,9 +30,6 @@
 from airflow.models.task_coordinates import TaskCoordinateResolver
 
 if TYPE_CHECKING:
-    from sqlalchemy.orm import Session
-
-    from airflow.models.dagbag import DBDagBag
     from airflow.models.task_coordinates import TaskCoordinate
 
 
@@ -63,8 +60,7 @@ def resolve_task_scope(
     dag_id: str,
     run_id: str,
     task_id: str,
-    session: Session,
-    dag_bag: DBDagBag,
+    resolver: TaskCoordinateResolver,
     map_index: int = -1,
     region_id: UUID | None = None,
     region_index: int | None = None,
@@ -83,7 +79,6 @@ def resolve_task_scope(
             map_index=region_index,
             region_id=region_id,
         )
-    resolver = TaskCoordinateResolver(dag_bag, session)
     if region_id == SENTINEL_REGION_ID or not resolver.has_regions(dag_id, run_id, task_id):
         return TaskScope(
             dag_id=dag_id,
@@ -137,8 +132,7 @@ def _task_scope(
         dag_id=dag_id,
         run_id=dag_run_id,
         task_id=task_id,
-        session=resolver.session,
-        dag_bag=resolver.dag_bag,
+        resolver=resolver,
         map_index=map_index,
         region_id=region_id,
         region_index=region_index,

airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py

@@ -237,8 +237,8 @@ def _patch_ti_validate_request(
         .where(TI.dag_id == dag_id, TI.run_id == dag_run_id, TI.task_id == task_id)
         .options(joinedload(TI.rendered_task_instance_fields))
     )
+    resolver = TaskCoordinateResolver(dag_bag, session)
     if body.region_id is None:
-        resolver = TaskCoordinateResolver(dag_bag, session)
         for version in session.scalars(
             select(TI.dag_version_id)
             .where(
@@ -254,8 +254,7 @@ def _patch_ti_validate_request(
         dag_id=dag_id,
         run_id=dag_run_id,
         task_id=task_id,
-        session=session,
-        dag_bag=dag_bag,
+        resolver=resolver,
         map_index=map_index if map_index is not None else -1,
         region_id=body.region_id,
         region_index=body.region_index,

airflow-core/src/airflow/models/xcom.py

@@ -267,7 +267,7 @@ def set(
         task_id: str,
         run_id: str,
         map_index: int = -1,
-        region_id: UUID | None = SENTINEL_REGION_ID,
+        region_id: UUID = SENTINEL_REGION_ID,
         serialize: bool = True,
         dag_result: bool = False,
         mapped_length: int | None = None,
@@ -282,6 +282,7 @@ def set(
         :param task_id: Task ID.
         :param run_id: DAG run ID for the task.
         :param map_index: Optional map index to assign XCom for a mapped task.
+        :param region_id: Region owning the task instance. Defaults to the top-level region.
         :param serialize: Optional parameter to specify if value should be serialized or not.
             The default is ``True``.
         :param mapped_length: Length of the value, if it can be used to expand a

airflow-core/src/airflow/ui/src/layouts/Details/Gantt/utils.test.ts

@@ -142,8 +142,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: null,
+          id: "try-1",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: null,
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: null,
           state: null,
@@ -165,8 +169,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: null,
+          id: "try-2",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: null,
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: "2024-03-14T10:00:00+00:00",
           state: "running",
@@ -210,8 +218,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: "2024-03-14T10:05:00+00:00",
+          id: "try-3",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: null,
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: "2024-03-14T10:00:00+00:00",
           state: "success",
@@ -235,8 +247,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: "2024-03-14T10:05:00+00:00",
+          id: "try-4",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: "2024-03-14T09:59:00+00:00",
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: "2024-03-14T09:58:00+00:00",
           start_date: "2024-03-14T10:00:00+00:00",
           state: "success",
@@ -260,8 +276,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: "2024-03-14T10:05:00+00:00",
+          id: "try-5",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: "2024-03-14T09:59:00+00:00",
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: "2024-03-14T09:58:00+00:00",
           start_date: "2024-03-14T10:00:00+00:00",
           state: "success",
@@ -295,8 +315,12 @@ describe("transformGanttData", () => {
         allTries: [
           {
             end_date: null,
+            id: "running-try",
             is_mapped: false,
+            map_index: -1,
             queued_dttm: "2024-03-14T09:59:00+00:00",
+            region_id: "00000000-0000-0000-0000-000000000000",
+            region_index: -1,
             scheduled_dttm: null,
             start_date: "2024-03-14T10:00:00+00:00",
             state: "running",
@@ -324,8 +348,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: "2024-03-14T10:05:00+00:00",
+          id: "try-6",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: "2024-03-14T09:59:00+00:00",
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: "2024-03-14T10:00:00+00:00",
           state: "success",
@@ -348,8 +376,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: "2024-03-14T10:05:00+00:00",
+          id: "try-7",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: null,
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: "2024-03-14T10:00:00+00:00",
           state: "success",
@@ -371,8 +403,12 @@ describe("transformGanttData", () => {
       allTries: [
         {
           end_date: "2024-03-14T10:05:00+00:00",
+          id: "try-8",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: null,
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: "2024-03-14T10:00:00+00:00",
           state: "failed",
@@ -382,8 +418,12 @@ describe("transformGanttData", () => {
         },
         {
           end_date: "2024-03-14T09:55:00+00:00",
+          id: "try-9",
           is_mapped: false,
+          map_index: -1,
           queued_dttm: null,
+          region_id: "00000000-0000-0000-0000-000000000000",
+          region_index: -1,
           scheduled_dttm: null,
           start_date: "2024-03-14T09:50:00+00:00",
           state: "failed",