Coverage for /home/airflow/.local/lib/python3.12/site-packages/airflow/api_fastapi/common/db/task_instances.py: 100%

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17 

18from __future__ import annotations 

19 

20from sqlalchemy import Select 

21from sqlalchemy.orm import contains_eager, joinedload 

22 

23from airflow.models import Base 

24from airflow.models.dag_version import DagVersion 

25from airflow.models.dagrun import DagRun 

26from airflow.models.taskinstance import TaskInstance 

27 

28 

29def eager_load_TI_and_TIH_for_validation( 

30 query: Select, 

31 orm_model: Base | None = None, 

32) -> Select: 

33 """ 

34 Add JOINs and eager-loading options for TaskInstanceResponse and TaskInstanceHistoryResponse. 

35 

36 Adds ``join(dag_run)`` and ``outerjoin(dag_version)`` to the query and 

37 configures ``contains_eager`` so SQLAlchemy reuses those joins for 

38 populating the related objects (dag_run, dag_model, dag_version, bundle). 

39 This keeps the join logic centralised, avoids duplicate JOINs that would 

40 otherwise occur when combining explicit joins with ``joinedload``, and 

41 ensures ORDER BY / WHERE clauses on DagRun columns resolve correctly. 

42 

43 :param query: The SELECT statement to augment. 

44 :param orm_model: The ORM model to load options for (defaults to TaskInstance). 

45 """ 

46 if orm_model is None: 

47 orm_model = TaskInstance 

48 

49 query = query.join(orm_model.dag_run).outerjoin(orm_model.dag_version) 

50 query = query.options( 

51 contains_eager(orm_model.dag_run).options(joinedload(DagRun.dag_model)), 

52 contains_eager(orm_model.dag_version).options(joinedload(DagVersion.bundle)), 

53 ) 

54 if orm_model is TaskInstance: 

55 query = query.options( 

56 joinedload(orm_model.task_instance_note), 

57 joinedload(orm_model.rendered_task_instance_fields), 

58 ) 

59 return query