Skip to content

Commit 91af86e

Browse files
Subham-KRLXjason810496
Subham-KRLX
and
jason810496
authored
Fix log task instance join (#59973)
* Fix incorrect Log and TaskInstance join query * Fix log task instance join foreign keys and remove stray newsfragment * Remove unused task variable assignments in test_log_task_instance_join_correctness * Fix CI failures and address reviewer comments: revert unrelated doc changes, clean up tests, add Mapped type hints to Log relationships, and run ruff formatting * Restore newsfragment and fix missing import in log.py * Fix reviewer feedback: add license, move imports, and fix newsfragments * Fix ruff static check failure in log.py * Flatten primaryjoin definition to fix formatting * Fix newsfragments per review: restore 58337 and remove 59965 * Remove unused import in log.py --------- Co-authored-by: Jason(Zhe-You) Liu <[email protected]>
1 parent 85842d4 commit 91af86e

2 files changed

Lines changed: 89 additions & 4 deletions

File tree

‎airflow-core/src/airflow/models/log.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from airflow.utils.sqlalchemy import UtcDateTime, mapped_column
2929

3030
if TYPE_CHECKING:
31+
from airflow.models.dag import DagModel
3132
from airflow.models.taskinstance import TaskInstance
3233
from airflow.models.taskinstancekey import TaskInstanceKey
3334

@@ -50,18 +51,18 @@ class Log(Base):
5051
extra: Mapped[str | None] = mapped_column(Text, nullable=True)
5152
try_number: Mapped[int | None] = mapped_column(Integer, nullable=True)
5253

53-
dag_model = relationship(
54+
dag_model: Mapped[DagModel | None] = relationship(
5455
"DagModel",
5556
viewonly=True,
5657
foreign_keys=[dag_id],
5758
primaryjoin="Log.dag_id == DagModel.dag_id",
5859
)
5960

60-
task_instance = relationship(
61+
task_instance: Mapped[TaskInstance | None] = relationship(
6162
"TaskInstance",
6263
viewonly=True,
63-
foreign_keys=[task_id],
64-
primaryjoin="Log.task_id == TaskInstance.task_id",
64+
foreign_keys=[dag_id, task_id, run_id, map_index],
65+
primaryjoin="and_(Log.dag_id == TaskInstance.dag_id, Log.task_id == TaskInstance.task_id, Log.run_id == TaskInstance.run_id, Log.map_index == TaskInstance.map_index)",
6566
lazy="noload",
6667
)
6768

‎airflow-core/tests/unit/models/test_log.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
from __future__ import annotations
19+
20+
import pytest
21+
from sqlalchemy import select
22+
from sqlalchemy.orm import joinedload
23+
24+
from airflow.models.log import Log
25+
from airflow.operators.empty import EmptyOperator
26+
from airflow.utils.state import TaskInstanceState
27+
28+
pytestmark = pytest.mark.db_test
29+
30+
31+
class TestLogTaskInstanceReproduction:
32+
def test_log_task_instance_join_correctness(self, dag_maker, session):
33+
# Create dag_1 with a task
34+
with dag_maker("dag_1", session=session):
35+
EmptyOperator(task_id="common_task_id")
36+
37+
dr1 = dag_maker.create_dagrun()
38+
ti1 = dr1.get_task_instance("common_task_id")
39+
ti1.state = TaskInstanceState.SUCCESS
40+
session.merge(ti1)
41+
session.commit()
42+
43+
# Create dag_2 with the SAME task_id
44+
with dag_maker("dag_2", session=session):
45+
EmptyOperator(task_id="common_task_id")
46+
47+
dr2 = dag_maker.create_dagrun()
48+
ti2 = dr2.get_task_instance("common_task_id")
49+
ti2.state = TaskInstanceState.FAILED
50+
session.merge(ti2)
51+
session.commit()
52+
53+
# Create a log entry specifically for dag_1's task instance
54+
log = Log(
55+
event="test_event",
56+
task_instance=ti1,
57+
)
58+
session.add(log)
59+
session.commit()
60+
61+
# Query with joinedload to trigger the relationship join
62+
63+
stmt = select(Log).where(Log.id == log.id).options(joinedload(Log.task_instance))
64+
loaded_log = session.scalar(stmt)
65+
66+
assert loaded_log.task_instance is not None
67+
assert loaded_log.task_instance.dag_id == "dag_1"
68+
assert loaded_log.task_instance.run_id == ti1.run_id
69+
70+
# Verify incorrect join for second dag
71+
log2 = Log(
72+
event="test_event_2",
73+
task_instance=ti2,
74+
)
75+
session.add(log2)
76+
session.commit()
77+
78+
stmt2 = select(Log).where(Log.id == log2.id).options(joinedload(Log.task_instance))
79+
loaded_log2 = session.scalar(stmt2)
80+
81+
# This should fail if the join is ambiguous and picks the first one (dag_1)
82+
assert loaded_log2.task_instance is not None
83+
assert loaded_log2.task_instance.dag_id == "dag_2"
84+
assert loaded_log2.task_instance.run_id == ti2.run_id

0 commit comments

Comments
 (0)