Repository: apache/airflow
File: airflow-core/src/airflow/models/taskinstance.py (Python)

Code before:
        .join(DeadlineAlertModel, Deadline.deadline_alert_id == DeadlineAlertModel.id)
        .where(
            Deadline.dagrun_id == dagrun.id,
            Deadline.missed == false(),
            DeadlineAlertModel.reference[ReferenceModels.REFERENCE_TYPE_FIELD].as_string()
            == ReferenceModels.DagRunQueuedAtDeadline.__name__,
        )
    ).all()

    if not results:
        return

    for deadline, deadline_alert in results:
        # We can't use evaluate_with() since the new queued_at is not written to the DB yet.
        deadline_interval = timedelta(seconds=deadline_alert.interval)
        new_deadline_time = new_queued_at + deadline_interval

        log.info(
            "Recalculating deadline %s for DagRun %s.%s: old=%s, new=%s",
            deadline.id,
            dagrun.dag_id,
            dagrun.run_id,
            deadline.deadline_time,
            new_deadline_time,
        )
        deadline.deadline_time = new_deadline_time


def clear_task_instances(
    tis: list[TaskInstance],
    session: Session,
    dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
    run_on_latest_version: bool = False,
    prevent_running_task: bool | None = None,
) -> None:
    """
    Clear a set of task instances, but make sure the running ones get killed.

    Also sets Dagrun's `state` to QUEUED and `start_date` to the time of execution.
    But only for finished DRs (SUCCESS and FAILED).
    Doesn't clear DR's `state` and `start_date`for running
    DRs (QUEUED and RUNNING) because clearing the state for already
    running DR is redundant and clearing `start_date` affects DR's duration.

    :param tis: a list of task instances
    :param session: current session
    :param dag_run_state: state to set finished DagRuns to.
        If set to False, DagRuns state will not be changed.
    :param run_on_latest_version: whether to run on latest serialized DAG and Bundle version

    :meta private:

Review comment (suggestion):
Suggestion: Maybe leave a comment saying that the session is (or rather, must be) flushed by the calling function? Looks strange just having this dangling.

Code after:
        .join(DeadlineAlertModel, Deadline.deadline_alert_id == DeadlineAlertModel.id)
        .where(
            Deadline.dagrun_id == dagrun.id,
            Deadline.missed == false(),
            DeadlineAlertModel.reference[ReferenceModels.REFERENCE_TYPE_FIELD].as_string()
            == ReferenceModels.DagRunQueuedAtDeadline.__name__,
        )
    ).all()

    if not results:
        return

    for deadline, deadline_alert in results:
        # We can't use evaluate_with() since the new queued_at is not written to the DB yet.
        deadline_interval = timedelta(seconds=deadline_alert.interval)
        new_deadline_time = new_queued_at + deadline_interval

        log.debug(
            "Recalculating deadline %s for DagRun %s.%s: old=%s, new=%s",
            deadline.id,
            dagrun.dag_id,
            dagrun.run_id,
            deadline.deadline_time,
            new_deadline_time,
        )
        deadline.deadline_time = new_deadline_time
    # Do not flush/commit here in order to keep the scheduler loop atomic.
    # These changes are committed by the calling function.


def clear_task_instances(
    tis: list[TaskInstance],
    session: Session,
    dag_run_state: DagRunState | Literal[False] = DagRunState.QUEUED,
    run_on_latest_version: bool = False,
    prevent_running_task: bool | None = None,
) -> None:
    """
    Clear a set of task instances, but make sure the running ones get killed.

    Also sets Dagrun's `state` to QUEUED and `start_date` to the time of execution.
    But only for finished DRs (SUCCESS and FAILED).
    Doesn't clear DR's `state` and `start_date`for running
    DRs (QUEUED and RUNNING) because clearing the state for already
    running DR is redundant and clearing `start_date` affects DR's duration.

    :param tis: a list of task instances
    :param session: current session
    :param dag_run_state: state to set finished DagRuns to.
        If set to False, DagRuns state will not be changed.
    :param run_on_latest_version: whether to run on latest serialized DAG and Bundle version