# 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()