Repository: apache/airflow
File: airflow-core/src/airflow/jobs/scheduler_job_runner.py (Python)

Code before:
                .options(alias_loader)
                .options(selectinload(TI.dag_version))
                .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
                .join(DM, TI.dag_id == DM.dag_id)
                .where(
                    TI.state.in_((TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING)),
                    TI.last_heartbeat_at < limit_dttm,
                )
                .where(TI.queued_by_job_id == self.job.id)
            )
        )
        if task_instances_without_heartbeats:
            self.log.warning(
                "Failing %s TIs without heartbeat after %s",
                len(task_instances_without_heartbeats),
                limit_dttm,
            )
        return list(task_instances_without_heartbeats)

    def _purge_task_instances_without_heartbeats(
        self, task_instances_without_heartbeats: list[TI], *, session: Session
    ) -> None:
        dag_id_to_team_name: dict[str, str | None] = {}
        if conf.getboolean("core", "multi_team"):
            unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats}
            dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)

        for ti in task_instances_without_heartbeats:
            task_instance_heartbeat_timeout_message_details = (
                self._generate_task_instance_heartbeat_timeout_message_details(ti)
            )
            if not ti.dag_version:
                # If old ti from Airflow 2 and dag_version is None, skip heartbeat timeout handling.
                self.log.warning(
                    "DAG Version not found for TaskInstance %s. Skipping heartbeat timeout handling.",
                    ti,
                )
                continue
            request = TaskCallbackRequest(
                filepath=ti.dag_model.relative_fileloc or "",
                bundle_name=ti.dag_version.bundle_name,
                bundle_version=ti.dag_run.bundle_version,
                ti=ti,
                msg=str(task_instance_heartbeat_timeout_message_details),
                context_from_server=TIRunContext(
                    dag_run=DRDataModel.model_validate(ti.dag_run, from_attributes=True),
                    max_tries=ti.max_tries,
                    variables=[],
                    connections=[],
                    xcom_keys_to_clear=[],
                ),

Review comment (suggestion):
```suggestion
        if conf.getboolean("core", "multi_team"):
            unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats}
            dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)
        else:
            dag_id_to_team_name = {}
```

Code after:
                .options(alias_loader)
                .options(selectinload(TI.dag_version))
                .with_hint(TI, "USE INDEX (ti_state)", dialect_name="mysql")
                .join(DM, TI.dag_id == DM.dag_id)
                .where(
                    TI.state.in_((TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING)),
                    TI.last_heartbeat_at < limit_dttm,
                )
                .where(TI.queued_by_job_id == self.job.id)
            )
        )
        if task_instances_without_heartbeats:
            self.log.warning(
                "Failing %s TIs without heartbeat after %s",
                len(task_instances_without_heartbeats),
                limit_dttm,
            )
        return list(task_instances_without_heartbeats)

    def _purge_task_instances_without_heartbeats(
        self, task_instances_without_heartbeats: list[TI], *, session: Session
    ) -> None:
        if conf.getboolean("core", "multi_team"):
            unique_dag_ids = {ti.dag_id for ti in task_instances_without_heartbeats}
            dag_id_to_team_name = self._get_team_names_for_dag_ids(unique_dag_ids, session)
        else:
            dag_id_to_team_name = {}

        for ti in task_instances_without_heartbeats:
            task_instance_heartbeat_timeout_message_details = (
                self._generate_task_instance_heartbeat_timeout_message_details(ti)
            )
            if not ti.dag_version:
                # If old ti from Airflow 2 and dag_version is None, skip heartbeat timeout handling.
                self.log.warning(
                    "DAG Version not found for TaskInstance %s. Skipping heartbeat timeout handling.",
                    ti,
                )
                continue
            request = TaskCallbackRequest(
                filepath=ti.dag_model.relative_fileloc or "",
                bundle_name=ti.dag_version.bundle_name,
                bundle_version=ti.dag_run.bundle_version,
                ti=ti,
                msg=str(task_instance_heartbeat_timeout_message_details),
                context_from_server=TIRunContext(
                    dag_run=DRDataModel.model_validate(ti.dag_run, from_attributes=True),
                    max_tries=ti.max_tries,
                    variables=[],
                    connections=[],
                    xcom_keys_to_clear=[],