diff --git a/airflow/cli/commands/dag_command.py b/airflow/cli/commands/dag_command.py index 13dd351edc6f7..92dde41146f62 100644 --- a/airflow/cli/commands/dag_command.py +++ b/airflow/cli/commands/dag_command.py @@ -139,7 +139,7 @@ def dag_backfill(args, dag=None): pool=args.pool, delay_on_limit_secs=args.delay_on_limit, verbose=args.verbose, - conf=run_conf, + conf_dict=run_conf, rerun_failed_tasks=args.rerun_failed_tasks, run_backwards=args.run_backwards ) diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py index 5deeadc61378e..c4cb65edcfb49 100644 --- a/airflow/models/baseoperator.py +++ b/airflow/models/baseoperator.py @@ -1334,6 +1334,9 @@ def get_serialized_fields(cls): return cls.__serialized_fields + def is_subdag(self) -> bool: # noqa + return False + def chain(*tasks: Union[BaseOperator, Sequence[BaseOperator]]): r""" diff --git a/airflow/models/dag.py b/airflow/models/dag.py index 137400693b0d2..e3decd90593a3 100644 --- a/airflow/models/dag.py +++ b/airflow/models/dag.py @@ -38,7 +38,7 @@ from sqlalchemy.orm.session import Session from airflow import settings, utils -from airflow.configuration import conf +from airflow.configuration import conf as conf_values from airflow.dag.base_dag import BaseDag from airflow.exceptions import AirflowException, DuplicateTaskIdFound, TaskNotFound from airflow.models.base import ID_LEN, Base @@ -74,15 +74,14 @@ def get_last_dagrun(dag_id, session, include_externally_triggered=False): Last dag run can be any type of run eg. scheduled or backfilled. Overridden DagRuns are ignored. """ - DR = DagRun - query = session.query(DR).filter(DR.dag_id == dag_id) + query = session.query(DagRun).filter(DagRun.dag_id == dag_id) if not include_externally_triggered: - query = query.filter(DR.external_trigger == False) # noqa pylint: disable=singleton-comparison - query = query.order_by(DR.execution_date.desc()) + query = query.filter(DagRun.external_trigger == False) # noqa pylint: disable=singleton-comparison + query = query.order_by(DagRun.execution_date.desc()) return query.first() -@functools.total_ordering +@functools.total_ordering # pylint: disable=too-many-instance-attributes,too-many-public-methods class DAG(BaseDag, LoggingMixin): """ A dag (directed acyclic graph) is a collection of tasks with directional @@ -209,35 +208,27 @@ class DAG(BaseDag, LoggingMixin): __serialized_fields: Optional[FrozenSet[str]] = None - def __init__( - self, - dag_id: str, - description: Optional[str] = None, - schedule_interval: Optional[ScheduleInterval] = timedelta(days=1), - start_date: Optional[datetime] = None, - end_date: Optional[datetime] = None, - full_filepath: Optional[str] = None, - template_searchpath: Optional[Union[str, Iterable[str]]] = None, - template_undefined: Type[jinja2.Undefined] = jinja2.Undefined, - user_defined_macros: Optional[Dict] = None, - user_defined_filters: Optional[Dict] = None, - default_args: Optional[Dict] = None, - concurrency: int = conf.getint('core', 'dag_concurrency'), - max_active_runs: int = conf.getint('core', 'max_active_runs_per_dag'), - dagrun_timeout: Optional[timedelta] = None, - sla_miss_callback: Optional[Callable] = None, - default_view: str = conf.get('webserver', 'dag_default_view').lower(), - orientation: str = conf.get('webserver', 'dag_orientation'), - catchup: bool = conf.getboolean('scheduler', 'catchup_by_default'), - on_success_callback: Optional[DagStateChangeCallback] = None, - on_failure_callback: Optional[DagStateChangeCallback] = None, - doc_md: Optional[str] = None, - params: Optional[Dict] = None, - access_control: Optional[Dict] = None, - is_paused_upon_creation: Optional[bool] = None, - jinja_environment_kwargs: Optional[Dict] = None, - tags: Optional[List[str]] = None - ): + def __init__(self, # pylint: disable=too-many-arguments,too-many-locals + dag_id: str, description: Optional[str] = None, + schedule_interval: Optional[ScheduleInterval] = timedelta(days=1), + start_date: Optional[datetime] = None, end_date: Optional[datetime] = None, + full_filepath: Optional[str] = None, + template_searchpath: Optional[Union[str, Iterable[str]]] = None, + template_undefined: Type[jinja2.Undefined] = jinja2.Undefined, + user_defined_macros: Optional[Dict] = None, user_defined_filters: Optional[Dict] = None, + default_args: Optional[Dict] = None, + concurrency: int = conf_values.getint('core', 'dag_concurrency'), + max_active_runs: int = conf_values.getint('core', 'max_active_runs_per_dag'), + dagrun_timeout: Optional[timedelta] = None, sla_miss_callback: Optional[Callable] = None, + default_view: str = conf_values.get('webserver', 'dag_default_view').lower(), + orientation: str = conf_values.get('webserver', 'dag_orientation'), + catchup: bool = conf_values.getboolean('scheduler', 'catchup_by_default'), + on_success_callback: Optional[DagStateChangeCallback] = None, + on_failure_callback: Optional[DagStateChangeCallback] = None, doc_md: Optional[str] = None, + params: Optional[Dict] = None, access_control: Optional[Dict] = None, + is_paused_upon_creation: Optional[bool] = None, + jinja_environment_kwargs: Optional[Dict] = None, tags: Optional[List[str]] = None): + super().__init__() self.user_defined_macros = user_defined_macros self.user_defined_filters = user_defined_filters self.default_args = copy.deepcopy(default_args or {}) @@ -258,7 +249,7 @@ def __init__( self._description = description # set file location to caller source path - back = sys._getframe().f_back + back = sys._getframe().f_back # noqa pylint: disable=protected-access self.fileloc = back.f_code.co_filename if back else "" self.task_dict: Dict[str, BaseOperator] = {} @@ -334,7 +325,7 @@ def __repr__(self): return "".format(self=self) def __eq__(self, other): - if (type(self) == type(other) and + if (type(self) == type(other) and # pylint: disable=unidiomatic-typecheck self.dag_id == other.dag_id): # Use getattr() instead of __dict__ as __dict__ doesn't return @@ -350,12 +341,12 @@ def __lt__(self, other): def __hash__(self): hash_components = [type(self)] - for c in self._comps: + for component in self._comps: # task_ids returns a list and lists can't be hashed - if c == 'task_ids': + if component == 'task_ids': val = tuple(self.task_dict.keys()) else: - val = getattr(self, c, None) + val = getattr(self, component, None) try: hash(val) hash_components.append(val) @@ -379,27 +370,27 @@ def date_range( num: Optional[int] = None, end_date: Optional[datetime] = timezone.utcnow(), ) -> List[datetime]: + """Returns date range for the interval.""" if num is not None: end_date = None return utils_date_range( start_date=start_date, end_date=end_date, num=num, delta=self.normalized_schedule_interval) - def is_fixed_time_schedule(self): + def is_fixed_time_schedule(self) -> bool: """ Figures out if the DAG schedule has a fixed time (e.g. 3 AM). :return: True if the schedule has a fixed time, False if not. """ now = datetime.now() - cron = croniter(self.normalized_schedule_interval, now) + cron = croniter(str(self.normalized_schedule_interval), now) start = cron.get_next(datetime) cron_next = cron.get_next(datetime) if cron_next.minute == start.minute and cron_next.hour == start.hour: return True - return False def following_schedule(self, dttm): @@ -424,11 +415,12 @@ def following_schedule(self, dttm): else: # absolute (e.g. 3 AM) naive = cron.get_next(datetime) - tz = pendulum.timezone(self.timezone.name) - following = timezone.make_aware(naive, tz) + time_zone = pendulum.timezone(self.timezone.name) + following = timezone.make_aware(naive, time_zone) return timezone.convert_to_utc(following) elif self.normalized_schedule_interval is not None: return timezone.convert_to_utc(dttm + self.normalized_schedule_interval) + return None def previous_schedule(self, dttm): """ @@ -452,11 +444,12 @@ def previous_schedule(self, dttm): else: # absolute (e.g. 3 AM) naive = cron.get_prev(datetime) - tz = pendulum.timezone(self.timezone.name) - previous = timezone.make_aware(naive, tz) + time_zone = pendulum.timezone(self.timezone.name) + previous = timezone.make_aware(naive, time_zone) return timezone.convert_to_utc(previous) elif self.normalized_schedule_interval is not None: return timezone.convert_to_utc(dttm - self.normalized_schedule_interval) + return None def get_run_dates(self, start_date, end_date=None): """ @@ -506,68 +499,86 @@ def normalize_schedule(self, dttm): @provide_session def get_last_dagrun(self, session=None, include_externally_triggered=False): + """ + Get last dagrun. + """ return get_last_dagrun(self.dag_id, session=session, include_externally_triggered=include_externally_triggered) @property def dag_id(self) -> str: + """Get Dag ID,""" return self._dag_id @dag_id.setter def dag_id(self, value: str) -> None: + """Set Dag ID""" self._dag_id = value @property def full_filepath(self) -> str: + """Get Full DAG filepath""" return self._full_filepath @full_filepath.setter def full_filepath(self, value) -> None: + """Set Full DAG filepath""" self._full_filepath = value @property def concurrency(self) -> int: + """Get Concurrency""" return self._concurrency @concurrency.setter def concurrency(self, value: int): + """Set Concurrency""" self._concurrency = value @property def access_control(self): + """Get Access Control""" return self._access_control @access_control.setter def access_control(self, value): + """Set Access Control""" self._access_control = value @property def description(self) -> Optional[str]: + """Get description""" return self._description @property def default_view(self) -> str: + """Get default_view""" return self._default_view @property - def pickle_id(self) -> Optional[int]: + def pickle_id(self) -> Optional[int]: # pylint: disable=invalid-overridden-method + """Get pickle_id""" return self._pickle_id @pickle_id.setter - def pickle_id(self, value: int) -> None: + def pickle_id(self, value: int) -> None: # pylint: disable=invalid-overridden-method + """Set pickle_id""" self._pickle_id = value @property def tasks(self) -> List[BaseOperator]: + """Get tasks""" return list(self.task_dict.values()) @tasks.setter def tasks(self, val): + """Cannot set tasks""" raise AttributeError( 'DAG.tasks can not be modified. Use dag.add_task() instead.') @property def task_ids(self) -> List[str]: + """Get task ids""" return list(self.task_dict.keys()) @property @@ -575,9 +586,9 @@ def filepath(self) -> str: """ File location of where the dag object is instantiated """ - fn = self.full_filepath.replace(settings.DAGS_FOLDER + '/', '') - fn = fn.replace(os.path.dirname(__file__) + '/', '') - return fn + filepath = self.full_filepath.replace(settings.DAGS_FOLDER + '/', '') + filepath = filepath.replace(os.path.dirname(__file__) + '/', '') + return filepath @property def folder(self) -> str: @@ -596,7 +607,8 @@ def owner(self) -> str: @property def allow_future_exec_dates(self) -> bool: - return conf.getboolean( + """Check if it is ok to trigger future execution dates""" + return conf_values.getboolean( 'scheduler', 'allow_trigger_in_future', fallback=False) and self.schedule_interval is None @@ -607,10 +619,9 @@ def get_concurrency_reached(self, session=None) -> bool: Returns a boolean indicating whether the concurrency limit for this DAG has been reached """ - TI = TaskInstance - qry = session.query(func.count(TI.task_id)).filter( - TI.dag_id == self.dag_id, - TI.state == State.RUNNING, + qry = session.query(func.count(TaskInstance.task_id)).filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.state == State.RUNNING, ) return qry.scalar() >= self.concurrency @@ -691,7 +702,7 @@ def handle_callback(self, dagrun, success=True, reason=None, session=None): context.update({'reason': reason}) try: callback(context) - except Exception: + except Exception: # noqa pylint: disable=broad-except self.log.exception("failed to invoke dag state update callback") Stats.incr("dag.callback_exceptions") @@ -723,7 +734,7 @@ def get_num_active_runs(self, external_trigger=None, session=None): query = (session .query(func.count()) .filter(DagRun.dag_id == self.dag_id) - .filter(DagRun.state == State.RUNNING)) + .filter(DagRun.state == State.RUNNING)) # pylint: disable=comparison-with-callable if external_trigger is not None: query = query.filter(DagRun.external_trigger == external_trigger) @@ -796,20 +807,20 @@ def subdags(self): Returns a list of the subdag objects associated to this DAG """ # Check SubDag for class but don't check class directly - from airflow.operators.subdag_operator import SubDagOperator subdag_lst = [] for task in self.tasks: - if (isinstance(task, SubDagOperator) or - # TODO remove in Airflow 2.0 - type(task).__name__ == 'SubDagOperator' or - task.task_type == 'SubDagOperator'): + # TODO remove in Airflow 2.0 + if type(task).__name__ == 'SubDagOperator' \ + or task.task_type == 'SubDagOperator' \ + or task.is_subdag(): subdag_lst.append(task.subdag) subdag_lst += task.subdag.subdags return subdag_lst def resolve_template_files(self): - for t in self.tasks: - t.resolve_template_files() + """Resolves template files.""" + for task in self.tasks: + task.resolve_template_files() def get_template_env(self) -> jinja2.Environment: """Build a Jinja2 environment.""" @@ -829,7 +840,8 @@ def get_template_env(self) -> jinja2.Environment: if self.jinja_environment_kwargs: jinja_env_options.update(self.jinja_environment_kwargs) - env = jinja2.Environment(**jinja_env_options) # type: ignore + env = jinja2.Environment( # noqa + **jinja_env_options) # type: ignore # Add any user defined items. Safe to edit globals as long as no templates are rendered yet. # http://jinja.pocoo.org/docs/2.10/api/#jinja2.Environment.globals @@ -851,12 +863,13 @@ def set_dependency(self, upstream_task_id, downstream_task_id): @provide_session def get_task_instances( self, start_date=None, end_date=None, state=None, session=None): + """Returns task instances.""" if not start_date: start_date = (timezone.utcnow() - timedelta(30)).date() start_date = timezone.make_aware( datetime.combine(start_date, datetime.min.time())) - tis = session.query(TaskInstance).filter( + task_instance_query_result = session.query(TaskInstance).filter( TaskInstance.dag_id == self.dag_id, TaskInstance.execution_date >= start_date, TaskInstance.task_id.in_([t.task_id for t in self.tasks]), @@ -864,26 +877,14 @@ def get_task_instances( # This allows allow_trigger_in_future config to take affect, rather than mandating exec_date <= UTC if end_date or not self.allow_future_exec_dates: end_date = end_date or timezone.utcnow() - tis = tis.filter(TaskInstance.execution_date <= end_date) + task_instance_query_result = task_instance_query_result.filter( + TaskInstance.execution_date <= end_date) - if state: - if isinstance(state, str): - tis = tis.filter(TaskInstance.state == state) - else: - # this is required to deal with NULL values - if None in state: - if all(x is None for x in state): - tis = tis.filter(TaskInstance.state.is_(None)) - else: - not_none_state = [s for s in state if s] - tis = tis.filter( - or_(TaskInstance.state.in_(not_none_state), - TaskInstance.state.is_(None)) - ) - else: - tis = tis.filter(TaskInstance.state.in_(state)) - tis = tis.order_by(TaskInstance.execution_date).all() - return tis + task_instance_query_result = TaskInstance.filter_task_instance_by_state_or_state_array( + state, + task_instance_query_result) + task_instance_query_result = task_instance_query_result.order_by(TaskInstance.execution_date).all() + return task_instance_query_result @property def roots(self) -> List[BaseOperator]: @@ -906,8 +907,6 @@ def topological_sort(self, include_subdag_tasks: bool = False): :param include_subdag_tasks: whether to include tasks in subdags, default to False :return: list of tasks in topological order """ - from airflow.operators.subdag_operator import SubDagOperator # Avoid circular import - # convert into an OrderedDict to speedup lookup while keeping order the same graph_unsorted = OrderedDict((task.task_id, task) for task in self.tasks) @@ -917,6 +916,17 @@ def topological_sort(self, include_subdag_tasks: bool = False): if len(self.tasks) == 0: return tuple(graph_sorted) + self.sort_graph_in_a_loop(graph_sorted, graph_unsorted, include_subdag_tasks) + + return tuple(graph_sorted) + + def sort_graph_in_a_loop( # pylint: disable=too-many-nested-blocks + self, + graph_sorted, + graph_unsorted, + include_subdag_tasks): + """Sorts graph in a loop until it is sorted.""" + # Run until the unsorted graph is empty. while graph_unsorted: # Go through each of the node/edges pairs in the unsorted @@ -941,15 +951,13 @@ def topological_sort(self, include_subdag_tasks: bool = False): acyclic = True del graph_unsorted[node.task_id] graph_sorted.append(node) - if include_subdag_tasks and isinstance(node, SubDagOperator): + if include_subdag_tasks and node.is_subdag(): graph_sorted.extend(node.subdag.topological_sort(include_subdag_tasks=True)) if not acyclic: raise AirflowException("A cyclic dependency occurred in dag: {}" .format(self.dag_id)) - return tuple(graph_sorted) - @provide_session def set_dag_runs_state( self, @@ -958,6 +966,7 @@ def set_dag_runs_state( start_date: Optional[datetime] = None, end_date: Optional[datetime] = None, ) -> None: + """Sets the state of DagRun""" query = session.query(DagRun).filter_by(dag_id=self.dag_id) if start_date: query = query.filter(DagRun.execution_date >= start_date) @@ -966,7 +975,7 @@ def set_dag_runs_state( query.update({DagRun.state: state}) @provide_session - def clear( + def clear( # pylint: disable=too-many-arguments,too-many-locals self, start_date=None, end_date=None, only_failed=False, only_running=False, @@ -1016,20 +1025,19 @@ def clear( :param dag_bag: The DagBag used to find the dags :type dag_bag: airflow.models.dagbag.DagBag """ - TI = TaskInstance - tis = session.query(TI) + tis = session.query(TaskInstance) if include_subdags: # Crafting the right filter for dag_id and task_ids combo conditions = [] for dag in self.subdags + [self]: conditions.append( - (TI.dag_id == dag.dag_id) & - TI.task_id.in_(dag.task_ids) + (TaskInstance.dag_id == dag.dag_id) & + TaskInstance.task_id.in_(dag.task_ids) ) tis = tis.filter(or_(*conditions)) else: - tis = session.query(TI).filter(TI.dag_id == self.dag_id) - tis = tis.filter(TI.task_id.in_(self.task_ids)) + tis = session.query(TaskInstance).filter(TaskInstance.dag_id == self.dag_id) + tis = tis.filter(TaskInstance.task_id.in_(self.task_ids)) if include_parentdag and self.is_subdag and self.parent_dag is not None: p_dag = self.parent_dag.sub_dag( @@ -1053,68 +1061,20 @@ def clear( )) if start_date: - tis = tis.filter(TI.execution_date >= start_date) + tis = tis.filter(TaskInstance.execution_date >= start_date) if end_date: - tis = tis.filter(TI.execution_date <= end_date) + tis = tis.filter(TaskInstance.execution_date <= end_date) if only_failed: tis = tis.filter(or_( - TI.state == State.FAILED, - TI.state == State.UPSTREAM_FAILED)) + TaskInstance.state == State.FAILED, + TaskInstance.state == State.UPSTREAM_FAILED)) if only_running: - tis = tis.filter(TI.state == State.RUNNING) + tis = tis.filter(TaskInstance.state == State.RUNNING) if include_subdags: - from airflow.sensors.external_task_sensor import ExternalTaskMarker - - # Recursively find external tasks indicated by ExternalTaskMarker - instances = tis.all() - for ti in instances: - if ti.operator == ExternalTaskMarker.__name__: - task: ExternalTaskMarker = cast(ExternalTaskMarker, copy.copy(self.get_task(ti.task_id))) - ti.task = task - - if recursion_depth == 0: - # Maximum recursion depth allowed is the recursion_depth of the first - # ExternalTaskMarker in the tasks to be cleared. - max_recursion_depth = task.recursion_depth - - if recursion_depth + 1 > max_recursion_depth: - # Prevent cycles or accidents. - raise AirflowException("Maximum recursion depth {} reached for {} {}. " - "Attempted to clear too many tasks " - "or there may be a cyclic dependency." - .format(max_recursion_depth, - ExternalTaskMarker.__name__, ti.task_id)) - ti.render_templates() - external_tis = session.query(TI).filter(TI.dag_id == task.external_dag_id, - TI.task_id == task.external_task_id, - TI.execution_date == - pendulum.parse(task.execution_date)) - - for tii in external_tis: - if not dag_bag: - dag_bag = DagBag() - external_dag = dag_bag.get_dag(tii.dag_id) - if not external_dag: - raise AirflowException("Could not find dag {}".format(tii.dag_id)) - downstream = external_dag.sub_dag( - task_regex=r"^{}$".format(tii.task_id), - include_upstream=False, - include_downstream=True - ) - tis = tis.union(downstream.clear(start_date=tii.execution_date, - end_date=tii.execution_date, - only_failed=only_failed, - only_running=only_running, - confirm_prompt=confirm_prompt, - include_subdags=include_subdags, - include_parentdag=False, - dag_run_state=dag_run_state, - get_tis=True, - session=session, - recursion_depth=recursion_depth + 1, - max_recursion_depth=max_recursion_depth, - dag_bag=dag_bag)) + tis = self._retrieve_subdags(confirm_prompt, dag_bag, dag_run_state, include_subdags, + max_recursion_depth, only_failed, only_running, recursion_depth, + session, tis) if get_tis: return tis @@ -1133,7 +1093,7 @@ def clear( return 0 if confirm_prompt: ti_list = "\n".join([str(t) for t in tis]) - question = ( + question = ( # noqa "You are about to delete these {count} tasks:\n" "{ti_list}\n\n" "Are you sure? (yes/no): ").format(count=count, ti_list=ti_list) @@ -1159,8 +1119,84 @@ def clear( session.commit() return count + def _retrieve_subdags( # pylint: disable=too-many-arguments + self, confirm_prompt, dag_bag, dag_run_state, include_subdags, max_recursion_depth, + only_failed, only_running, recursion_depth, session, tis): + from airflow.sensors.external_task_sensor import ExternalTaskMarker + + # Recursively find external tasks indicated by ExternalTaskMarker + instances = tis.all() + for ti in instances: + if ti.operator == ExternalTaskMarker.__name__: + task: ExternalTaskMarker = cast(ExternalTaskMarker, copy.copy(self.get_task(ti.task_id))) + ti.task = task + + if recursion_depth == 0: + # Maximum recursion depth allowed is the recursion_depth of the first + # ExternalTaskMarker in the tasks to be cleared. + max_recursion_depth = task.recursion_depth + + if recursion_depth + 1 > max_recursion_depth: + # Prevent cycles or accidents. + raise AirflowException("Maximum recursion depth {} reached for {} {}. " + "Attempted to clear too many tasks " + "or there may be a cyclic dependency." + .format(max_recursion_depth, + ExternalTaskMarker.__name__, ti.task_id)) + ti.render_templates() + tis = self._process_external_dags(confirm_prompt, dag_bag, dag_run_state, include_subdags, + max_recursion_depth, only_failed, only_running, + recursion_depth, session, task, tis) + return tis + + @staticmethod + def _process_external_dags( # pylint: disable=too-many-arguments + confirm_prompt, + dag_bag, + dag_run_state, + include_subdags, + max_recursion_depth, + only_failed, + only_running, + recursion_depth, + session, + task, + tis + ): + external_tis = session.query(TaskInstance).filter( + TaskInstance.dag_id == task.external_dag_id, + TaskInstance.task_id == task.external_task_id, + TaskInstance.execution_date == + pendulum.parse(task.execution_date)) + for tii in external_tis: + if not dag_bag: + dag_bag = DagBag() + external_dag = dag_bag.get_dag(tii.dag_id) + if not external_dag: + raise AirflowException("Could not find dag {}".format(tii.dag_id)) + downstream = external_dag.sub_dag( + task_regex=r"^{}$".format(tii.task_id), + include_upstream=False, + include_downstream=True + ) + tis = tis.union(downstream.clear( + start_date=tii.execution_date, + end_date=tii.execution_date, + only_failed=only_failed, + only_running=only_running, + confirm_prompt=confirm_prompt, + include_subdags=include_subdags, + include_parentdag=False, + dag_run_state=dag_run_state, + get_tis=True, + session=session, + recursion_depth=recursion_depth + 1, + max_recursion_depth=max_recursion_depth, + dag_bag=dag_bag)) + return tis + @classmethod - def clear_dags( + def clear_dags( # pylint: disable=too-many-arguments cls, dags, start_date=None, end_date=None, @@ -1172,6 +1208,7 @@ def clear_dags( dag_run_state=State.RUNNING, dry_run=False, ): + """Clears DAGs.""" all_tis = [] for dag in dags: tis = dag.clear( @@ -1196,7 +1233,7 @@ def clear_dags( return 0 if confirm_prompt: ti_list = "\n".join([str(t) for t in all_tis]) - question = ( + question = ( # noqa "You are about to delete these {} tasks:\n" "{}\n\n" "Are you sure? (yes/no): ").format(count, ti_list) @@ -1219,14 +1256,13 @@ def clear_dags( return count def __deepcopy__(self, memo): - # Swiwtcharoo to go around deepcopying objects coming through the - # backdoor + # Switch to go around deep-copying objects coming through the backdoor cls = self.__class__ result = cls.__new__(cls) memo[id(self)] = result for k, v in self.__dict__.items(): if k not in ('user_defined_macros', 'user_defined_filters', 'params'): - setattr(result, k, copy.deepcopy(v, memo)) + setattr(result, k, copy.deepcopy(v, memo)) # noqa result.user_defined_macros = self.user_defined_macros result.user_defined_filters = self.user_defined_filters @@ -1251,21 +1287,24 @@ def sub_dag(self, task_regex, include_downstream=False, regex_match = [ t for t in self.tasks if re.findall(task_regex, t.task_id)] also_include = [] - for t in regex_match: + for task in regex_match: if include_downstream: - also_include += t.get_flat_relatives(upstream=False) + also_include += task.get_flat_relatives(upstream=False) if include_upstream: - also_include += t.get_flat_relatives(upstream=True) + also_include += task.get_flat_relatives(upstream=True) # Compiling the unique list of tasks that made the cut # Make sure to not recursively deepcopy the dag while copying the task - dag.task_dict = {t.task_id: copy.deepcopy(t, {id(t.dag): dag}) + dag.task_dict = {t.task_id: copy.deepcopy(t, {id(t.dag): dag}) # noqa for t in regex_match + also_include} - for t in dag.tasks: + for task in dag.tasks: # Removing upstream/downstream references to tasks that did not # made the cut - t._upstream_task_ids = t.upstream_task_ids.intersection(dag.task_dict.keys()) - t._downstream_task_ids = t.downstream_task_ids.intersection( + # pylint: disable=protected-access + task._upstream_task_ids = \ + task.upstream_task_ids.intersection(dag.task_dict.keys()) + # pylint: disable=protected-access + task._downstream_task_ids = task.downstream_task_ids.intersection( dag.task_dict.keys()) if len(dag.tasks) < len(self.tasks): @@ -1274,9 +1313,11 @@ def sub_dag(self, task_regex, include_downstream=False, return dag def has_task(self, task_id: str): - return task_id in (t.task_id for t in self.tasks) + """Whether the task belongs to the Dag.""" + return task_id in (task.task_id for task in self.tasks) def get_task(self, task_id: str, include_subdags: bool = False) -> BaseOperator: + """Retrieve operator for the task id.""" if task_id in self.task_dict: return self.task_dict[task_id] if include_subdags: @@ -1286,49 +1327,51 @@ def get_task(self, task_id: str, include_subdags: bool = False) -> BaseOperator: raise TaskNotFound("Task {task_id} not found".format(task_id=task_id)) def pickle_info(self): - d = {} - d['is_picklable'] = True + """Return pickle info.""" + pickle_dict = {'is_picklable': True} try: dttm = timezone.utcnow() pickled = pickle.dumps(self) - d['pickle_len'] = len(pickled) - d['pickling_duration'] = "{}".format(timezone.utcnow() - dttm) - except Exception as e: + pickle_dict['pickle_len'] = len(pickled) + pickle_dict['pickling_duration'] = "{}".format(timezone.utcnow() - dttm) + except Exception as e: # pylint: disable=broad-except self.log.debug(e) - d['is_picklable'] = False - d['stacktrace'] = traceback.format_exc() - return d + pickle_dict['is_picklable'] = False + pickle_dict['stacktrace'] = traceback.format_exc() + return pickle_dict @provide_session def pickle(self, session=None) -> DagPickle: + """Pickle DAG.""" dag = session.query( DagModel).filter(DagModel.dag_id == self.dag_id).first() - dp = None + dag_pickle = None if dag and dag.pickle_id: - dp = session.query(DagPickle).filter( + dag_pickle = session.query(DagPickle).filter( DagPickle.id == dag.pickle_id).first() - if not dp or dp.pickle != self: - dp = DagPickle(dag=self) - session.add(dp) - self.last_pickled = timezone.utcnow() + if not dag_pickle or dag_pickle.pickle != self: + dag_pickle = DagPickle(dag=self) + session.add(dag_pickle) + self.last_pickled = timezone.utcnow() # noqa pylint: disable=attribute-defined-outside-init session.commit() - self.pickle_id = dp.id + self.pickle_id = dag_pickle.id - return dp + return dag_pickle def tree_view(self) -> None: """Print an ASCII tree representation of the DAG.""" def get_downstream(task, level=0): print((" " * level * 4) + str(task)) level += 1 - for t in task.downstream_list: - get_downstream(t, level) + for downstream_task in task.downstream_list: + get_downstream(downstream_task, level) - for t in self.roots: - get_downstream(t) + for _task in self.roots: + get_downstream(_task) @property def task(self): + """Return task for the DAG.""" from airflow.operators.python import task return functools.partial(task, dag=self) @@ -1364,7 +1407,7 @@ def add_task(self, task): self.task_dict[task.task_id] = task task.dag = self - self.task_count = len(self.task_dict) + self.task_count = len(self.task_dict) # noqa pylint: disable=attribute-defined-outside-init def add_tasks(self, tasks): """ @@ -1376,20 +1419,20 @@ def add_tasks(self, tasks): for task in tasks: self.add_task(task) - def run( + def run( # pylint: disable=too-many-arguments self, start_date=None, end_date=None, mark_success=False, local=False, executor=None, - donot_pickle=conf.getboolean('core', 'donot_pickle'), + donot_pickle=conf_values.getboolean('core', 'donot_pickle'), ignore_task_deps=False, ignore_first_depends_on_past=True, pool=None, delay_on_limit_secs=1.0, verbose=False, - conf=None, + conf_dict=None, rerun_failed_tasks=False, run_backwards=False, ): @@ -1420,8 +1463,8 @@ def run( :type delay_on_limit_secs: float :param verbose: Make logging output more verbose :type verbose: bool - :param conf: user defined dictionary passed from CLI - :type conf: dict + :param conf_dict: user defined dictionary passed from CLI + :type conf_dict: dict :param rerun_failed_tasks: :type: bool :param run_backwards: @@ -1447,7 +1490,7 @@ def run( pool=pool, delay_on_limit_secs=delay_on_limit_secs, verbose=verbose, - conf=conf, + conf=conf_dict, rerun_failed_tasks=rerun_failed_tasks, run_backwards=run_backwards, ) @@ -1530,7 +1573,10 @@ def create_dagrun(self, @classmethod @provide_session - def bulk_sync_to_db(cls, dags: Collection["DAG"], sync_time=None, session=None): + def bulk_sync_to_db(cls, + dags: Collection["DAG"], # pylint: disable=unsubscriptable-object + sync_time=None, + session=None): """ Save attributes about list of DAG to the DB. Note that this method can be called for both DAGs and SubDAGs. A SubDag is actually a @@ -1540,6 +1586,7 @@ def bulk_sync_to_db(cls, dags: Collection["DAG"], sync_time=None, session=None): :type dags: List[airflow.models.dag.DAG] :param sync_time: The time that the DAG should be marked as sync'ed :type sync_time: datetime + :param session: DB Session :return: None """ if not dags: @@ -1572,13 +1619,25 @@ def bulk_sync_to_db(cls, dags: Collection["DAG"], sync_time=None, session=None): session.add(orm_dag) orm_dags.append(orm_dag) + cls._update_orm_dags_with_dag(dag_by_ids, orm_dags, session, sync_time) + + if settings.STORE_DAG_CODE: + DagCode.bulk_sync_to_db([dag.fileloc for dag in orm_dags]) + + session.commit() + + for dag in dags: + cls.bulk_sync_to_db(dag.subdags, sync_time=sync_time, session=session) + + @classmethod + def _update_orm_dags_with_dag(cls, dag_by_ids, orm_dags, session, sync_time): for orm_dag in sorted(orm_dags, key=lambda d: d.dag_id): dag = dag_by_ids[orm_dag.dag_id] if dag.is_subdag: orm_dag.is_subdag = True - orm_dag.fileloc = dag.parent_dag.fileloc # type: ignore - orm_dag.root_dag_id = dag.parent_dag.dag_id # type: ignore - orm_dag.owners = dag.parent_dag.owner # type: ignore + orm_dag.fileloc = dag.parent_dag.fileloc # noqa type: ignore + orm_dag.root_dag_id = dag.parent_dag.dag_id # noqa type: ignore + orm_dag.owners = dag.parent_dag.owner # noqa type: ignore else: orm_dag.is_subdag = False orm_dag.fileloc = dag.fileloc @@ -1588,25 +1647,21 @@ def bulk_sync_to_db(cls, dags: Collection["DAG"], sync_time=None, session=None): orm_dag.default_view = dag.default_view orm_dag.description = dag.description orm_dag.schedule_interval = dag.schedule_interval - for orm_tag in list(orm_dag.tags): - if orm_tag.name not in orm_dag.tags: - session.delete(orm_tag) - orm_dag.tags.remove(orm_tag) - if dag.tags: - orm_tag_names = [t.name for t in orm_dag.tags] - for dag_tag in list(dag.tags): - if dag_tag not in orm_tag_names: - dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id) - orm_dag.tags.append(dag_tag_orm) - session.add(dag_tag_orm) - - if settings.STORE_DAG_CODE: - DagCode.bulk_sync_to_db([dag.fileloc for dag in orm_dags]) + cls._update_tags_in_orm_dags(dag, orm_dag, session) - session.commit() - - for dag in dags: - cls.bulk_sync_to_db(dag.subdags, sync_time=sync_time, session=session) + @classmethod + def _update_tags_in_orm_dags(cls, dag, orm_dag, session): + for orm_tag in list(orm_dag.tags): + if orm_tag.name not in orm_dag.tags: + session.delete(orm_tag) + orm_dag.tags.remove(orm_tag) + if dag.tags: + orm_tag_names = [t.name for t in orm_dag.tags] + for dag_tag in list(dag.tags): + if dag_tag not in orm_tag_names: + dag_tag_orm = DagTag(name=dag_tag, dag_id=dag.dag_id) + orm_dag.tags.append(dag_tag_orm) + session.add(dag_tag_orm) @provide_session def sync_to_db(self, sync_time=None, session=None): @@ -1617,6 +1672,7 @@ def sync_to_db(self, sync_time=None, session=None): :param sync_time: The time that the DAG should be marked as sync'ed :type sync_time: datetime + :param session: DB session :return: None """ self.bulk_sync_to_db([self], sync_time, session) @@ -1630,6 +1686,7 @@ def deactivate_unknown_dags(active_dag_ids, session=None): :param active_dag_ids: list of DAG IDs that are active :type active_dag_ids: list[unicode] + :param session: DB session :return: None """ @@ -1651,6 +1708,7 @@ def deactivate_stale_dags(expiration_date, session=None): :param expiration_date: set inactive DAGs that were touched before this time :type expiration_date: datetime + :param session: DB session :return: None """ for dag in session.query( @@ -1677,29 +1735,23 @@ def get_num_task_instances(dag_id, task_ids=None, states=None, session=None): :type task_ids: list[unicode] :param states: A list of states to filter by if supplied :type states: list[state] + :param session: DB session :return: The number of running tasks :rtype: int """ - qry = session.query(func.count(TaskInstance.task_id)).filter( + task_instance_query_result = session.query(func.count(TaskInstance.task_id)).filter( TaskInstance.dag_id == dag_id, ) if task_ids: - qry = qry.filter( + task_instance_query_result = task_instance_query_result.filter( TaskInstance.task_id.in_(task_ids), ) if states: - if None in states: - if all(x is None for x in states): - qry = qry.filter(TaskInstance.state.is_(None)) - else: - not_none_states = [state for state in states if state] - qry = qry.filter(or_( - TaskInstance.state.in_(not_none_states), - TaskInstance.state.is_(None))) - else: - qry = qry.filter(TaskInstance.state.in_(states)) - return qry.scalar() + task_instance_query_result = TaskInstance.filter_task_instance_by_state_array( + states, + task_instance_query_result) + return task_instance_query_result.scalar() @classmethod def get_serialized_fields(cls): @@ -1725,7 +1777,7 @@ class DagTag(Base): dag_id = Column(String(ID_LEN), ForeignKey('dag.dag_id'), primary_key=True) def __repr__(self): - return self.name + return str(self.name) class DagModel(Base): @@ -1739,7 +1791,7 @@ class DagModel(Base): root_dag_id = Column(String(ID_LEN)) # A DAG can be paused from the UI / DB # Set this default value of is_paused based on a configuration value! - is_paused_at_creation = conf\ + is_paused_at_creation = conf_values\ .getboolean('core', 'dags_are_paused_at_creation') is_paused = Column(Boolean, default=is_paused_at_creation) @@ -1783,20 +1835,24 @@ def __repr__(self): @property def timezone(self): + """Return timezone.""" return settings.TIMEZONE @staticmethod @provide_session def get_dagmodel(dag_id, session=None): + """Return DagModel from the DB.""" return session.query(DagModel).filter(DagModel.dag_id == dag_id).first() @classmethod @provide_session def get_current(cls, dag_id, session=None): + """Get current DagModel as class method from the DB.""" return session.query(cls).filter(cls.dag_id == dag_id).first() @provide_session def get_last_dagrun(self, session=None, include_externally_triggered=False): + """Get last dagrun for the Dag.""" return get_last_dagrun(self.dag_id, session=session, include_externally_triggered=include_externally_triggered) @@ -1812,7 +1868,7 @@ def get_paused_dag_ids(dag_ids: List[str], session: Session = None) -> Set[str]: """ paused_dag_ids = ( session.query(DagModel.dag_id) - .filter(DagModel.is_paused.is_(True)) + .filter(DagModel.is_paused.is_(True)) # noqa .filter(DagModel.dag_id.in_(dag_ids)) .all() ) @@ -1822,6 +1878,7 @@ def get_paused_dag_ids(dag_ids: List[str], session: Session = None) -> Set[str]: @property def safe_dag_id(self): + """Safe dag it (with . replaced with __dot__)""" return self.dag_id.replace('.', '__dot__') @provide_session @@ -1865,22 +1922,31 @@ def deactivate_deleted_dags(cls, alive_dag_filelocs: List[str], session=None): dag_models = session.query(cls).all() try: - for dag_model in dag_models: - if dag_model.fileloc is not None: - if correct_maybe_zipped(dag_model.fileloc) not in alive_dag_filelocs: - dag_model.is_active = False - else: - # If is_active is set as False and the DAG File still exists - # Change is_active=True - if not dag_model.is_active: - dag_model.is_active = True - else: - continue + cls._deactivate_deleted_dags_in_models(alive_dag_filelocs, dag_models) session.commit() except Exception: session.rollback() raise + @classmethod + def _deactivate_deleted_dags_in_models(cls, alive_dag_filelocs, dag_models): + """Loop through all models and deactivate DAGs.""" + for dag_model in dag_models: + if dag_model.fileloc is not None: + cls._check_if_model_is_active(alive_dag_filelocs, dag_model) + else: + continue + + @classmethod + def _check_if_model_is_active(cls, alive_dag_filelocs, dag_model): + if correct_maybe_zipped(dag_model.fileloc) not in alive_dag_filelocs: + dag_model.is_active = False + else: + # If is_active is set as False and the DAG File still exists + # Change is_active=True + if not dag_model.is_active: + dag_model.is_active = True + class DagContext: """ @@ -1907,12 +1973,14 @@ class DagContext: @classmethod def push_context_managed_dag(cls, dag: DAG): + """Push context to stack""" if cls._context_managed_dag: cls._previous_context_managed_dags.append(cls._context_managed_dag) cls._context_managed_dag = dag @classmethod def pop_context_managed_dag(cls) -> Optional[DAG]: + """Pop context to stack""" old_dag = cls._context_managed_dag if cls._previous_context_managed_dags: cls._context_managed_dag = cls._previous_context_managed_dags.pop() @@ -1922,4 +1990,5 @@ def pop_context_managed_dag(cls) -> Optional[DAG]: @classmethod def get_current_dag(cls) -> Optional[DAG]: + """Get the current context.""" return cls._context_managed_dag diff --git a/airflow/models/dagrun.py b/airflow/models/dagrun.py index 736da96f018db..8df2b51fb621d 100644 --- a/airflow/models/dagrun.py +++ b/airflow/models/dagrun.py @@ -19,7 +19,7 @@ from typing import Any, List, Optional, Tuple, Union from sqlalchemy import ( - Boolean, Column, DateTime, Index, Integer, PickleType, String, UniqueConstraint, and_, func, or_, + Boolean, Column, DateTime, Index, Integer, PickleType, String, UniqueConstraint, and_, func, ) from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.declarative import declared_attr @@ -28,7 +28,7 @@ from airflow.exceptions import AirflowException from airflow.models.base import ID_LEN, Base -from airflow.models.taskinstance import TaskInstance as TI +from airflow.models.taskinstance import TaskInstance from airflow.settings import task_instance_mutation_hook from airflow.stats import Stats from airflow.ti_deps.dep_context import DepContext @@ -89,7 +89,7 @@ def __init__( super().__init__() def __repr__(self): - return ( + return ( # noqa '' ).format( @@ -99,15 +99,19 @@ def __repr__(self): external_trigger=self.external_trigger) def get_state(self): + """Returns DagRun state""" return self._state def set_state(self, state): + """Sets DagRun state.""" + if self._state != state: self._state = state self.end_date = timezone.utcnow() if self._state in State.finished() else None @declared_attr def state(self): + """Attribute for state""" return synonym('_state', descriptor=property(self.get_state, self.set_state)) @provide_session @@ -118,14 +122,12 @@ def refresh_from_db(self, session: Session = None): :param session: database session :type session: Session """ - DR = DagRun - exec_date = func.cast(self.execution_date, DateTime) - dr = session.query(DR).filter( - DR.dag_id == self.dag_id, - func.cast(DR.execution_date, DateTime) == exec_date, - DR.run_id == self.run_id + dr = session.query(DagRun).filter( + DagRun.dag_id == self.dag_id, + func.cast(DagRun.execution_date, DateTime) == exec_date, + DagRun.run_id == self.run_id ).one() self.id = dr.id @@ -170,35 +172,33 @@ def find( :param execution_end_date: dag run that was executed until this date :type execution_end_date: datetime.datetime """ - DR = DagRun - - qry = session.query(DR) + qry = session.query(DagRun) dag_ids = [dag_id] if isinstance(dag_id, str) else dag_id if dag_ids: - qry = qry.filter(DR.dag_id.in_(dag_ids)) + qry = qry.filter(DagRun.dag_id.in_(dag_ids)) if run_id: - qry = qry.filter(DR.run_id == run_id) + qry = qry.filter(DagRun.run_id == run_id) if execution_date: if isinstance(execution_date, list): - qry = qry.filter(DR.execution_date.in_(execution_date)) + qry = qry.filter(DagRun.execution_date.in_(execution_date)) else: - qry = qry.filter(DR.execution_date == execution_date) + qry = qry.filter(DagRun.execution_date == execution_date) if execution_start_date and execution_end_date: - qry = qry.filter(DR.execution_date.between(execution_start_date, execution_end_date)) + qry = qry.filter(DagRun.execution_date.between(execution_start_date, execution_end_date)) elif execution_start_date: - qry = qry.filter(DR.execution_date >= execution_start_date) + qry = qry.filter(DagRun.execution_date >= execution_start_date) elif execution_end_date: - qry = qry.filter(DR.execution_date <= execution_end_date) + qry = qry.filter(DagRun.execution_date <= execution_end_date) if state: - qry = qry.filter(DR.state == state) + qry = qry.filter(DagRun.state == state) # pylint: disable=comparison-with-callable if external_trigger is not None: - qry = qry.filter(DR.external_trigger == external_trigger) + qry = qry.filter(DagRun.external_trigger == external_trigger) if run_type: - qry = qry.filter(DR.run_type == run_type.value) + qry = qry.filter(DagRun.run_type == run_type.value) if no_backfills: - qry = qry.filter(DR.run_type != DagRunType.BACKFILL_JOB.value) + qry = qry.filter(DagRun.run_type != DagRunType.BACKFILL_JOB.value) - dr = qry.order_by(DR.execution_date).all() + dr = qry.order_by(DagRun.execution_date).all() return dr @@ -207,36 +207,26 @@ def generate_run_id(run_type: DagRunType, execution_date: datetime) -> str: """Generate Run ID based on Run Type and Execution Date""" return f"{run_type.value}__{execution_date.isoformat()}" + # pylint: disable=no-member + @provide_session def get_task_instances(self, state=None, session=None): """ Returns the task instances for this dag run """ - tis = session.query(TI).filter( - TI.dag_id == self.dag_id, - TI.execution_date == self.execution_date, + task_instance_query_result = session.query(TaskInstance).filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.execution_date == self.execution_date, ) - if state: - if isinstance(state, str): - tis = tis.filter(TI.state == state) - else: - # this is required to deal with NULL values - if None in state: - if all(x is None for x in state): - tis = tis.filter(TI.state.is_(None)) - else: - not_none_state = [s for s in state if s] - tis = tis.filter( - or_(TI.state.in_(not_none_state), - TI.state.is_(None)) - ) - else: - tis = tis.filter(TI.state.in_(state)) + task_instance_query_result = TaskInstance.filter_task_instance_by_state_or_state_array( + state, + task_instance_query_result) if self.dag and self.dag.partial: - tis = tis.filter(TI.task_id.in_(self.dag.task_ids)) - return tis.all() + task_instance_query_result = task_instance_query_result.filter( + TaskInstance.task_id.in_(self.dag.task_ids)) # pylint: disable=no-member + return task_instance_query_result.all() @provide_session def get_task_instance(self, task_id: str, session: Session = None): @@ -248,10 +238,10 @@ def get_task_instance(self, task_id: str, session: Session = None): :param session: Sqlalchemy ORM Session :type session: Session """ - ti = session.query(TI).filter( - TI.dag_id == self.dag_id, - TI.execution_date == self.execution_date, - TI.task_id == task_id + ti = session.query(TaskInstance).filter( + TaskInstance.dag_id == self.dag_id, + TaskInstance.execution_date == self.execution_date, + TaskInstance.task_id == task_id ).first() return ti @@ -276,7 +266,7 @@ def get_previous_dagrun(self, state: Optional[str] = None, session: Session = No DagRun.execution_date < self.execution_date, ] if state is not None: - filters.append(DagRun.state == state) + filters.append(DagRun.state == state) # pylint: disable=comparison-with-callable return session.query(DagRun).filter( *filters ).order_by( @@ -294,7 +284,7 @@ def get_previous_scheduled_dagrun(self, session: Session = None): ).first() @provide_session - def update_state(self, session: Session = None) -> List[TI]: + def update_state(self, session: Session = None) -> List[TaskInstance]: """ Determines the overall state of the DagRun based on the state of its TaskInstances. @@ -306,7 +296,7 @@ def update_state(self, session: Session = None) -> List[TI]: """ dag = self.get_dag() - ready_tis: List[TI] = [] + ready_tis: List[TaskInstance] = [] tis = list(self.get_task_instances(session=session, state=State.task_states + (State.SHUTDOWN,))) self.log.debug("number of tis tasks for %s: %s task(s)", self, len(tis)) for ti in tis: @@ -318,6 +308,7 @@ def update_state(self, session: Session = None) -> List[TI]: none_depends_on_past = all(not t.task.depends_on_past for t in unfinished_tasks) none_task_concurrency = all(t.task.task_concurrency is None for t in unfinished_tasks) + are_runnable_tasks = False if unfinished_tasks: scheduleable_tasks = [ut for ut in unfinished_tasks if ut.state in SCHEDULEABLE_STATES] self.log.debug( @@ -327,7 +318,7 @@ def update_state(self, session: Session = None) -> List[TI]: self.log.debug("ready tis length for %s: %s task(s)", self, len(ready_tis)) if none_depends_on_past and none_task_concurrency: # small speed up - are_runnable_tasks = ready_tis or self._are_premature_tis( + are_runnable_tasks = len(ready_tis) > 0 or self._are_premature_tis( unfinished_tasks, finished_tasks, session) or changed_tis duration = (timezone.utcnow() - start_dttm) @@ -344,7 +335,7 @@ def update_state(self, session: Session = None) -> List[TI]: self.set_state(State.FAILED) dag.handle_callback(self, success=False, reason='task_failure', session=session) - # if all leafs succeeded and no unfinished tasks, the run succeeded + # if all leaves succeeded and no unfinished tasks, the run succeeded elif not unfinished_tasks and all( leaf_ti.state in {State.SUCCESS, State.SKIPPED} for leaf_ti in leaf_tis ): @@ -372,49 +363,49 @@ def update_state(self, session: Session = None) -> List[TI]: return ready_tis + @staticmethod def _get_ready_tis( - self, - scheduleable_tasks: List[TI], - finished_tasks: List[TI], + scheduleable_tasks: List[TaskInstance], + finished_tasks: List[TaskInstance], session: Session, - ) -> Tuple[List[TI], bool]: + ) -> Tuple[List[TaskInstance], bool]: old_states = {} - ready_tis: List[TI] = [] + ready_tis: List[TaskInstance] = [] changed_tis = False if not scheduleable_tasks: return ready_tis, changed_tis # Check dependencies - for st in scheduleable_tasks: - old_state = st.state - if st.are_dependencies_met( + for scheduleable_task in scheduleable_tasks: + old_state = scheduleable_task.state + if scheduleable_task.are_dependencies_met( dep_context=DepContext( flag_upstream_failed=True, finished_tasks=finished_tasks), session=session): - ready_tis.append(st) + ready_tis.append(scheduleable_task) else: - old_states[st.key] = old_state + old_states[scheduleable_task.key] = old_state # Check if any ti changed state - tis_filter = TI.filter_for_tis(old_states.keys()) + tis_filter = TaskInstance.filter_for_tis(old_states.keys()) if tis_filter is not None: - fresh_tis = session.query(TI).filter(tis_filter).all() + fresh_tis = session.query(TaskInstance).filter(tis_filter).all() changed_tis = any(ti.state != old_states[ti.key] for ti in fresh_tis) return ready_tis, changed_tis + @staticmethod def _are_premature_tis( - self, - unfinished_tasks: List[TI], - finished_tasks: List[TI], + unfinished_tasks: List[TaskInstance], + finished_tasks: List[TaskInstance], session: Session, ) -> bool: # there might be runnable tasks that are up for retry and for some reason(retry delay, etc) are # not ready yet so we set the flags to count them in - for ut in unfinished_tasks: - if ut.are_dependencies_met( + for unfinished_task in unfinished_tasks: + if unfinished_task.are_dependencies_met( dep_context=DepContext( flag_upstream_failed=True, ignore_in_retry_period=True, @@ -481,7 +472,7 @@ def verify_integrity(self, session: Session = None): Stats.incr( "task_instance_created-{}".format(task.__class__.__name__), 1, 1) - ti = TI(task, self.execution_date) + ti = TaskInstance(task, self.execution_date) task_instance_mutation_hook(ti) session.add(ti) @@ -489,11 +480,13 @@ def verify_integrity(self, session: Session = None): session.commit() except IntegrityError as err: self.log.info(str(err)) - self.log.info('Hit IntegrityError while creating the TIs for ' - f'{dag.dag_id} - {self.execution_date}.') + self.log.info('Hit IntegrityError while creating the TIs for %s - %s.', + dag.dag_id, self.execution_date) self.log.info('Doing session rollback.') session.rollback() + # pylint: enable=no-member + @staticmethod def get_run(session: Session, dag_id: str, execution_date: datetime): """ @@ -518,6 +511,7 @@ def get_run(session: Session, dag_id: str, execution_date: datetime): @property def is_backfill(self): + """Is this a backfill job?""" return self.run_type == DagRunType.BACKFILL_JOB.value @classmethod diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index e0c465e9ed19d..66a83788ab8a9 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -1773,6 +1773,33 @@ def filter_for_tis( raise TypeError("All elements must have the same type: `TaskInstance` or `TaskInstanceKey`.") + @staticmethod + def filter_task_instance_by_state_or_state_array(state, tis): + """Filter task instance by state or state array.""" + if state: + if isinstance(state, str): + tis = tis.filter(TaskInstance.state == state) + else: + tis = TaskInstance.filter_task_instance_by_state_array(state, tis) + return tis + + @staticmethod + def filter_task_instance_by_state_array(state, tasks_instance_query_result): + """Filter Task Instance by array of states - dealing with None values in the array.""" + if None in state: + if all(x is None for x in state): + tasks_instance_query_result = \ + tasks_instance_query_result.filter(TaskInstance.state.is_(None)) # noqa + else: + not_none_state = [s for s in state if s] + tasks_instance_query_result = tasks_instance_query_result.filter( + or_(TaskInstance.state.in_(not_none_state), + TaskInstance.state.is_(None)) # noqa + ) + else: + tasks_instance_query_result = tasks_instance_query_result.filter(TaskInstance.state.in_(state)) + return tasks_instance_query_result + # State of the task instance. # Stores string version of the task state. diff --git a/airflow/operators/subdag_operator.py b/airflow/operators/subdag_operator.py index 660687069ac08..ec05e3c6a4d7a 100644 --- a/airflow/operators/subdag_operator.py +++ b/airflow/operators/subdag_operator.py @@ -212,3 +212,6 @@ def _skip_downstream_tasks(self, context): self.skip(context['dag_run'], context['execution_date'], downstream_tasks) self.log.info('Done.') + + def is_subdag(self) -> bool: + return True diff --git a/scripts/ci/pylint_todo.txt b/scripts/ci/pylint_todo.txt index ac9fdc2fa5203..e69de29bb2d1d 100644 --- a/scripts/ci/pylint_todo.txt +++ b/scripts/ci/pylint_todo.txt @@ -1,3 +0,0 @@ -./airflow/models/dag.py -./airflow/models/dagrun.py -./airflow/www/utils.py diff --git a/tests/www/test_views.py b/tests/www/test_views.py index cd7f08dae461b..67f461c33e428 100644 --- a/tests/www/test_views.py +++ b/tests/www/test_views.py @@ -48,6 +48,7 @@ from airflow.jobs.base_job import BaseJob from airflow.models import DAG, Connection, DagRun, TaskInstance from airflow.models.baseoperator import BaseOperator, BaseOperatorLink +from airflow.models.dag import DagModel from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF from airflow.models.serialized_dag import SerializedDagModel from airflow.operators.bash import BashOperator @@ -988,8 +989,7 @@ def test_delete_dag_button_for_dag_on_scheduler_only(self): dag_id = 'example_bash_operator' test_dag_id = "non_existent_dag" - DM = models.DagModel - dag_query = self.session.query(DM).filter(DM.dag_id == dag_id) + dag_query = self.session.query(DagModel).filter(DagModel.dag_id == dag_id) dag_query.first().tags = [] # To avoid "FOREIGN KEY constraint" error self.session.commit() @@ -1000,7 +1000,7 @@ def test_delete_dag_button_for_dag_on_scheduler_only(self): self.check_content_in_response('/delete?dag_id={}'.format(test_dag_id), resp) self.check_content_in_response("return confirmDeleteDag(this, '{}')".format(test_dag_id), resp) - self.session.query(DM).filter(DM.dag_id == test_dag_id).update({'dag_id': dag_id}) + self.session.query(DagModel).filter(DagModel.dag_id == test_dag_id).update({'dag_id': dag_id}) self.session.commit() @parameterized.expand(["graph", "tree"]) @@ -1017,7 +1017,7 @@ def test_show_external_log_redirect_link_with_local_log_handler(self, endpoint): @mock.patch('airflow.utils.log.log_reader.TaskLogReader.log_handler', new_callable=PropertyMock) def test_show_external_log_redirect_link_with_external_log_handler(self, endpoint, mock_log_handler): """Show external links if log handler is external.""" - class ExternalHandler(ExternalLoggingMixin): + class ExternalHandler(ExternalLoggingMixin): # noqa LOG_NAME = 'ExternalLog' @property @@ -1100,7 +1100,7 @@ def setUp(self): with conf_vars({('logging', 'logging_config_class'): 'airflow_local_settings.LOGGING_CONFIG'}): self.app = application.create_app(testing=True) - self.appbuilder = self.app.appbuilder # pylint: disable=no-member + self.appbuilder = self.app.appbuilder # noqa pylint: disable=no-member self.app.config['WTF_CSRF_ENABLED'] = False self.client = self.app.test_client() settings.configure_orm() @@ -1172,7 +1172,7 @@ def test_get_logs_with_metadata_as_download_file(self): "task_id={}&execution_date={}&" \ "try_number={}&metadata={}&format=file" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), try_number, @@ -1200,7 +1200,7 @@ def test_get_logs_with_metadata_as_download_large_file(self): "task_id={}&execution_date={}&" \ "try_number={}&metadata={}&format=file" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), try_number, @@ -1217,7 +1217,7 @@ def test_get_logs_with_metadata(self): "task_id={}&execution_date={}&" \ "try_number={}&metadata={}" response = \ - self.client.get(url_template.format(self.DAG_ID, + self.client.get(url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), 1, @@ -1226,17 +1226,17 @@ def test_get_logs_with_metadata(self): password='test'), follow_redirects=True) - self.assertIn('"message":', response.data.decode('utf-8')) - self.assertIn('"metadata":', response.data.decode('utf-8')) - self.assertIn('Log for testing.', response.data.decode('utf-8')) - self.assertEqual(200, response.status_code) + self.assertIn('"message":', response.data.decode('utf-8')) # noqa + self.assertIn('"metadata":', response.data.decode('utf-8')) # noqa + self.assertIn('Log for testing.', response.data.decode('utf-8')) # noqa + self.assertEqual(200, response.status_code) # noqa def test_get_logs_with_null_metadata(self): url_template = "get_logs_with_metadata?dag_id={}&" \ "task_id={}&execution_date={}&" \ "try_number={}&metadata=null" response = \ - self.client.get(url_template.format(self.DAG_ID, + self.client.get(url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), 1), data=dict( @@ -1255,7 +1255,7 @@ def test_get_logs_with_metadata_for_removed_dag(self, mock_read): url_template = "get_logs_with_metadata?dag_id={}&" \ "task_id={}&execution_date={}&" \ "try_number={}&metadata={}" - url = url_template.format(self.DAG_ID_REMOVED, self.TASK_ID, + url = url_template.format(self.DAG_ID_REMOVED, self.TASK_ID, # noqa quote_plus(self.DEFAULT_DATE.isoformat()), 1, json.dumps({})) response = self.client.get( url, @@ -1276,7 +1276,7 @@ def test_get_logs_response_with_ti_equal_to_none(self): "task_id={}&execution_date={}&" \ "try_number={}&metadata={}&format=file" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa 'Non_Existing_ID', quote_plus(self.DEFAULT_DATE.isoformat()), try_number, @@ -1292,7 +1292,7 @@ def test_get_logs_with_json_response_format(self): "task_id={}&execution_date={}&" \ "try_number={}&metadata={}&format=json" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), try_number, @@ -1311,7 +1311,7 @@ def test_get_logs_for_handler_without_read_method(self, mock_log_reader): "task_id={}&execution_date={}&" \ "try_number={}&metadata={}&format=json" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), try_number, @@ -1334,7 +1334,7 @@ def test_redirect_to_external_log_with_local_log_handler(self, task_id): "task_id={}&execution_date={}&" \ "try_number={}" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa task_id, quote_plus(self.DEFAULT_DATE.isoformat()), try_number) @@ -1357,7 +1357,7 @@ def get_external_log_url(self, *args, **kwargs): "task_id={}&execution_date={}&" \ "try_number={}" try_number = 1 - url = url_template.format(self.DAG_ID, + url = url_template.format(self.DAG_ID, # noqa self.TASK_ID, quote_plus(self.DEFAULT_DATE.isoformat()), try_number) @@ -1649,7 +1649,7 @@ class TestDagACLView(TestBase): def setUpClass(cls): super().setUpClass() dagbag = models.DagBag(include_examples=True) - DAG.bulk_sync_to_db(dagbag.dags.values()) + DAG.bulk_sync_to_db(dagbag.dags.values()) # noqa def prepare_dagruns(self): dagbag = models.DagBag(include_examples=True) @@ -2303,7 +2303,7 @@ class TestTaskInstanceView(TestBase): TI_ENDPOINT = '/taskinstance/list/?_flt_0_execution_date={}' def test_start_date_filter(self): - resp = self.client.get(self.TI_ENDPOINT.format( + resp = self.client.get(self.TI_ENDPOINT.format( # noqa self.percent_encode('2018-10-09 22:44:31'))) # We aren't checking the logic of the date filter itself (that is built # in to FAB) but simply that our UTC conversion was run - i.e. it @@ -2315,7 +2315,7 @@ class TestTaskRescheduleView(TestBase): TI_ENDPOINT = '/taskreschedule/list/?_flt_0_execution_date={}' def test_start_date_filter(self): - resp = self.client.get(self.TI_ENDPOINT.format( + resp = self.client.get(self.TI_ENDPOINT.format( # noqa self.percent_encode('2018-10-09 22:44:31'))) # We aren't checking the logic of the date filter itself (that is built # in to FAB) but simply that our UTC conversion was run - i.e. it @@ -2416,7 +2416,7 @@ class TestTriggerDag(TestBase): def setUp(self): super().setUp() - self.session = Session() + self.session = Session() # noqa models.DagBag().get_dag("example_bash_operator").sync_to_db(session=self.session) def test_trigger_dag_button_normal_exist(self): @@ -2429,13 +2429,12 @@ def test_trigger_dag_button(self): test_dag_id = "example_bash_operator" - DR = models.DagRun - self.session.query(DR).delete() + self.session.query(DagRun).delete() self.session.commit() self.client.post('trigger?dag_id={}'.format(test_dag_id)) - run = self.session.query(DR).filter(DR.dag_id == test_dag_id).first() + run = self.session.query(DagRun).filter(DagRun.dag_id == test_dag_id).first() self.assertIsNotNone(run) self.assertIn(DagRunType.MANUAL.value, run.run_id) self.assertEqual(run.run_type, DagRunType.MANUAL.value) @@ -2446,13 +2445,12 @@ def test_trigger_dag_conf(self): test_dag_id = "example_bash_operator" conf_dict = {'string': 'Hello, World!'} - DR = models.DagRun - self.session.query(DR).delete() + self.session.query(DagRun).delete() self.session.commit() self.client.post('trigger?dag_id={}'.format(test_dag_id), data={'conf': json.dumps(conf_dict)}) - run = self.session.query(DR).filter(DR.dag_id == test_dag_id).first() + run = self.session.query(DagRun).filter(DagRun.dag_id == test_dag_id).first() self.assertIsNotNone(run) self.assertIn(DagRunType.MANUAL.value, run.run_id) self.assertEqual(run.run_type, DagRunType.MANUAL.value) @@ -2461,14 +2459,13 @@ def test_trigger_dag_conf(self): def test_trigger_dag_conf_malformed(self): test_dag_id = "example_bash_operator" - DR = models.DagRun - self.session.query(DR).delete() + self.session.query(DagRun).delete() self.session.commit() response = self.client.post('trigger?dag_id={}'.format(test_dag_id), data={'conf': '{"a": "b"'}) self.check_content_in_response('Invalid JSON configuration', response) - run = self.session.query(DR).filter(DR.dag_id == test_dag_id).first() + run = self.session.query(DagRun).filter(DagRun.dag_id == test_dag_id).first() self.assertIsNone(run) def test_trigger_dag_form(self): @@ -2533,7 +2530,7 @@ class AirflowLink(BaseOperatorLink): def get_link(self, operator, dttm): return 'https://airflow.apache.org' - class DummyTestOperator(BaseOperator): + class DummyTestOperator(BaseOperator): # noqa operator_extra_links = ( RaiseErrorLink(),