Some Pylint fixes in airflow/models/taskinstance.py (#9674)

apache · Jul 6, 2020 · 631ac48 · 631ac48
1 parent e764ea5
commit 631ac48
Show file tree

Hide file tree

Showing 8 changed files with 218 additions and 132 deletions.
diff --git a/airflow/executors/base_executor.py b/airflow/executors/base_executor.py
@@ -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

diff --git a/airflow/jobs/base_job.py b/airflow/jobs/base_job.py
@@ -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
@@ -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)
@@ -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: