Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions airflow/executors/base_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,7 @@
from typing import Any, Dict, List, Optional, Set, Tuple, Union

from airflow.configuration import conf
from airflow.models import TaskInstance
from airflow.models.taskinstance import SimpleTaskInstance, TaskInstanceKeyType
from airflow.models.taskinstance import SimpleTaskInstance, TaskInstance, TaskInstanceKeyType
from airflow.stats import Stats
from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.state import State
Expand Down
25 changes: 12 additions & 13 deletions airflow/jobs/base_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,12 @@
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm.session import make_transient

from airflow import models
from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.executors.executor_loader import ExecutorLoader
from airflow.models.base import ID_LEN, Base
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
from airflow.stats import Stats
from airflow.utils import helpers, timezone
from airflow.utils.helpers import convert_camel_to_snake
Expand Down Expand Up @@ -268,22 +269,20 @@ def reset_state_for_orphaned_tasks(self, filter_by_dag_run=None, session=None):
running_tis = self.executor.running

resettable_states = [State.SCHEDULED, State.QUEUED]
TI = models.TaskInstance
DR = models.DagRun
if filter_by_dag_run is None:
resettable_tis = (
session
.query(TI)
.query(TaskInstance)
.join(
DR,
DagRun,
and_(
TI.dag_id == DR.dag_id,
TI.execution_date == DR.execution_date))
TaskInstance.dag_id == DagRun.dag_id,
TaskInstance.execution_date == DagRun.execution_date))
.filter(
# pylint: disable=comparison-with-callable
DR.state == State.RUNNING,
DR.run_type != DagRunType.BACKFILL_JOB.value,
TI.state.in_(resettable_states))).all()
DagRun.state == State.RUNNING,
DagRun.run_type != DagRunType.BACKFILL_JOB.value,
TaskInstance.state.in_(resettable_states))).all()
else:
resettable_tis = filter_by_dag_run.get_task_instances(state=resettable_states,
session=session)
Expand All @@ -300,9 +299,9 @@ def query(result, items):
if not items:
return result

filter_for_tis = TI.filter_for_tis(items)
reset_tis = session.query(TI).filter(
filter_for_tis, TI.state.in_(resettable_states)
filter_for_tis = TaskInstance.filter_for_tis(items)
reset_tis = session.query(TaskInstance).filter(
filter_for_tis, TaskInstance.state.in_(resettable_states)
).with_for_update().all()

for ti in reset_tis:
Expand Down
Loading