Source code for airflow.models.trigger

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import datetime
from typing import Any, Dict, List, Optional

from sqlalchemy import Column, Integer, String, func

from airflow.models.base import Base
from airflow.models.taskinstance import TaskInstance
from airflow.triggers.base import BaseTrigger
from airflow.utils import timezone
from airflow.utils.session import provide_session
from airflow.utils.sqlalchemy import ExtendedJSON, UtcDateTime
from airflow.utils.state import State


[docs]class Trigger(Base): """ Triggers are a workload that run in an asynchronous event loop shared with other Triggers, and fire off events that will unpause deferred Tasks, start linked DAGs, etc. They are persisted into the database and then re-hydrated into a "triggerer" process, where many are run at once. We model it so that there is a many-to-one relationship between Task and Trigger, for future deduplication logic to use. Rows will be evicted from the database when the triggerer detects no active Tasks/DAGs using them. Events are not stored in the database; when an Event is fired, the triggerer will directly push its data to the appropriate Task/DAG. """
[docs] __tablename__ = "trigger"
[docs] id = Column(Integer, primary_key=True)
[docs] classpath = Column(String(1000), nullable=False)
[docs] kwargs = Column(ExtendedJSON, nullable=False)
[docs] created_date = Column(UtcDateTime, nullable=False)
[docs] triggerer_id = Column(Integer, nullable=True)
def __init__( self, classpath: str, kwargs: Dict[str, Any], created_date: Optional[datetime.datetime] = None ): super().__init__() self.classpath = classpath self.kwargs = kwargs self.created_date = created_date or timezone.utcnow() @classmethod
[docs] def from_object(cls, trigger: BaseTrigger): """ Alternative constructor that creates a trigger row based directly off of a Trigger object. """ classpath, kwargs = trigger.serialize() return cls(classpath=classpath, kwargs=kwargs)
@classmethod @provide_session
[docs] def bulk_fetch(cls, ids: List[int], session=None) -> Dict[int, "Trigger"]: """ Fetches all of the Triggers by ID and returns a dict mapping ID -> Trigger instance """ return {obj.id: obj for obj in session.query(cls).filter(cls.id.in_(ids)).all()}
@classmethod @provide_session
[docs] def clean_unused(cls, session=None): """ Deletes all triggers that have no tasks/DAGs dependent on them (triggers have a one-to-many relationship to both) """ # Update all task instances with trigger IDs that are not DEFERRED to remove them session.query(TaskInstance).filter( TaskInstance.state != State.DEFERRED, TaskInstance.trigger_id.isnot(None) ).update({TaskInstance.trigger_id: None}) # Get all triggers that have no task instances depending on them... ids = [ trigger_id for (trigger_id,) in ( session.query(cls.id) .join(TaskInstance, cls.id == TaskInstance.trigger_id, isouter=True) .group_by(cls.id) .having(func.count(TaskInstance.trigger_id) == 0) ) ] # ...and delete them (we can't do this in one query due to MySQL) session.query(Trigger).filter(Trigger.id.in_(ids)).delete(synchronize_session=False)
@classmethod @provide_session
[docs] def submit_event(cls, trigger_id, event, session=None): """ Takes an event from an instance of itself, and triggers all dependent tasks to resume. """ for task_instance in session.query(TaskInstance).filter( TaskInstance.trigger_id == trigger_id, TaskInstance.state == State.DEFERRED ): # Add the event's payload into the kwargs for the task next_kwargs = task_instance.next_kwargs or {} next_kwargs["event"] = event.payload task_instance.next_kwargs = next_kwargs # Remove ourselves as its trigger task_instance.trigger_id = None # Finally, mark it as scheduled so it gets re-queued task_instance.state = State.SCHEDULED
@classmethod @provide_session
[docs] def submit_failure(cls, trigger_id, session=None): """ Called when a trigger has failed unexpectedly, and we need to mark everything that depended on it as failed. Notably, we have to actually run the failure code from a worker as it may have linked callbacks, so hilariously we have to re-schedule the task instances to a worker just so they can then fail. We use a special __fail__ value for next_method to achieve this that the runtime code understands as immediate-fail, and pack the error into next_kwargs. TODO: Once we have shifted callback (and email) handling to run on workers as first-class concepts, we can run the failure code here in-process, but we can't do that right now. """ for task_instance in session.query(TaskInstance).filter( TaskInstance.trigger_id == trigger_id, TaskInstance.state == State.DEFERRED ): # Add the error and set the next_method to the fail state task_instance.next_method = "__fail__" task_instance.next_kwargs = {"error": "Trigger failure"} # Remove ourselves as its trigger task_instance.trigger_id = None # Finally, mark it as scheduled so it gets re-queued task_instance.state = State.SCHEDULED
@classmethod @provide_session
[docs] def ids_for_triggerer(cls, triggerer_id, session=None): """Retrieves a list of triggerer_ids.""" return [row[0] for row in session.query(cls.id).filter(cls.triggerer_id == triggerer_id)]
@classmethod @provide_session
[docs] def assign_unassigned(cls, triggerer_id, capacity, session=None): """ Takes a triggerer_id and the capacity for that triggerer and assigns unassigned triggers until that capacity is reached, or there are no more unassigned triggers. """ from airflow.jobs.base_job import BaseJob # To avoid circular import count = session.query(cls.id).filter(cls.triggerer_id == triggerer_id).count() capacity -= count if capacity <= 0: return alive_triggerer_ids = [ row[0] for row in session.query(BaseJob.id).filter( BaseJob.end_date is None, BaseJob.latest_heartbeat > timezone.utcnow() - datetime.timedelta(seconds=30), BaseJob.job_type == "TriggererJob", ) ] # Find triggers who do NOT have an alive triggerer_id, and then assign # up to `capacity` of those to us. trigger_ids_query = ( session.query(cls.id).filter(cls.triggerer_id.notin_(alive_triggerer_ids)).limit(capacity).all() ) session.query(cls).filter(cls.id.in_([i.id for i in trigger_ids_query])).update( {cls.triggerer_id: triggerer_id}, synchronize_session=False, ) session.commit()

Was this entry helpful?