From a827175fa352297def4bea82f066333357eced68 Mon Sep 17 00:00:00 2001
From: Jake Kramer <jkramer@ginkgobioworks.com>
Date: Wed, 23 Feb 2022 14:25:21 -0500
Subject: [PATCH 1/3] Fixes assignment of unassigned triggers

Previously, the query returned no alive triggerers which resulted
in all triggers to be assigned to the current triggerer. This works
fine, despite the logic bug, in the case where there's a single
triggerer. But with multiple triggerers, concurrent iterations of
the TriggerJob loop would bounce trigger ownership to whichever
loop ran last.

Addresses https://github.com/apache/airflow/issues/21616
---
 airflow/models/trigger.py    |  9 ++++---
 tests/models/test_trigger.py | 52 ++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/airflow/models/trigger.py b/airflow/models/trigger.py
index fba739ce00c05..6ac87ce89fd95 100644
--- a/airflow/models/trigger.py
+++ b/airflow/models/trigger.py
@@ -17,7 +17,7 @@
 import datetime
 from typing import Any, Dict, Iterable, Optional
 
-from sqlalchemy import Column, Integer, String, func
+from sqlalchemy import Column, Integer, String, func, or_
 
 from airflow.models.base import Base
 from airflow.models.taskinstance import TaskInstance
@@ -175,7 +175,7 @@ def assign_unassigned(cls, triggerer_id, capacity, session=None):
         alive_triggerer_ids = [
             row[0]
             for row in session.query(BaseJob.id).filter(
-                BaseJob.end_date is None,
+                BaseJob.end_date == None,  # noqa: E711
                 BaseJob.latest_heartbeat > timezone.utcnow() - datetime.timedelta(seconds=30),
                 BaseJob.job_type == "TriggererJob",
             )
@@ -184,7 +184,10 @@ def assign_unassigned(cls, triggerer_id, capacity, session=None):
         # 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.id)
+            .filter(or_(cls.triggerer_id.notin_(alive_triggerer_ids), cls.triggerer_id == None))  # noqa: E711
+            .limit(capacity)
+            .all()
         )
         session.query(cls).filter(cls.id.in_([i.id for i in trigger_ids_query])).update(
             {cls.triggerer_id: triggerer_id},
diff --git a/tests/models/test_trigger.py b/tests/models/test_trigger.py
index aacfa8859e150..99cd71f7679da 100644
--- a/tests/models/test_trigger.py
+++ b/tests/models/test_trigger.py
@@ -15,8 +15,11 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import datetime
+
 import pytest
 
+from airflow.jobs.triggerer_job import TriggererJob
 from airflow.models import TaskInstance, Trigger
 from airflow.operators.dummy import DummyOperator
 from airflow.triggers.base import TriggerEvent
@@ -36,9 +39,11 @@ def session():
 def clear_db(session):
     session.query(TaskInstance).delete()
     session.query(Trigger).delete()
+    session.query(TriggererJob).delete()
     yield session
     session.query(TaskInstance).delete()
     session.query(Trigger).delete()
+    session.query(TriggererJob).delete()
     session.commit()
 
 
@@ -124,3 +129,50 @@ def test_submit_failure(session, create_task_instance):
     updated_task_instance = session.query(TaskInstance).one()
     assert updated_task_instance.state == State.SCHEDULED
     assert updated_task_instance.next_method == "__fail__"
+
+
+def test_assign_unassigned(session, create_task_instance):
+    """
+    Tests that unassigned triggers of all appropriate states are assigned.
+    """
+    finished_triggerer = TriggererJob(None, heartrate=10, state=State.SUCCESS)
+    finished_triggerer.end_date = timezone.utcnow() - datetime.timedelta(hours=1)
+    session.add(finished_triggerer)
+    assert not finished_triggerer.is_alive()
+    healthy_triggerer = TriggererJob(None, heartrate=10, state=State.RUNNING)
+    session.add(healthy_triggerer)
+    assert healthy_triggerer.is_alive()
+    new_triggerer = TriggererJob(None, heartrate=10, state=State.RUNNING)
+    session.add(new_triggerer)
+    assert new_triggerer.is_alive()
+    session.commit()
+    trigger_on_healthy_triggerer = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})
+    trigger_on_healthy_triggerer.id = 1
+    trigger_on_healthy_triggerer.triggerer_id = healthy_triggerer.id
+    trigger_on_killed_triggerer = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})
+    trigger_on_killed_triggerer.id = 2
+    trigger_on_killed_triggerer.triggerer_id = finished_triggerer.id
+    trigger_unassigned_to_triggerer = Trigger(classpath="airflow.triggers.testing.SuccessTrigger", kwargs={})
+    trigger_unassigned_to_triggerer.id = 3
+    assert trigger_unassigned_to_triggerer.triggerer_id is None
+    session.add(trigger_on_healthy_triggerer)
+    session.add(trigger_on_killed_triggerer)
+    session.add(trigger_unassigned_to_triggerer)
+    session.commit()
+    assert session.query(Trigger).count() == 3
+    Trigger.assign_unassigned(new_triggerer.id, 100, session=session)
+    session.expire_all()
+    # Check that trigger on killed triggerer and unassigned trigger are assigned to new triggerer
+    assert (
+        session.query(Trigger).filter(Trigger.id == trigger_on_killed_triggerer.id).one().triggerer_id
+        == new_triggerer.id
+    )
+    assert (
+        session.query(Trigger).filter(Trigger.id == trigger_unassigned_to_triggerer.id).one().triggerer_id
+        == new_triggerer.id
+    )
+    # Check that trigger on healthy triggerer still assigned to existing triggerer
+    assert (
+        session.query(Trigger).filter(Trigger.id == trigger_on_healthy_triggerer.id).one().triggerer_id
+        == healthy_triggerer.id
+    )

From 11d19dfbb68fba5dd945c53e38fa1358d63b9bac Mon Sep 17 00:00:00 2001
From: Jake Kramer <jkramer@ginkgobioworks.com>
Date: Wed, 23 Feb 2022 15:34:27 -0500
Subject: [PATCH 2/3] Address comments

---
 airflow/models/trigger.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/airflow/models/trigger.py b/airflow/models/trigger.py
index 6ac87ce89fd95..f394cc55b487e 100644
--- a/airflow/models/trigger.py
+++ b/airflow/models/trigger.py
@@ -175,7 +175,7 @@ def assign_unassigned(cls, triggerer_id, capacity, session=None):
         alive_triggerer_ids = [
             row[0]
             for row in session.query(BaseJob.id).filter(
-                BaseJob.end_date == None,  # noqa: E711
+                BaseJob.end_date.is_(None),
                 BaseJob.latest_heartbeat > timezone.utcnow() - datetime.timedelta(seconds=30),
                 BaseJob.job_type == "TriggererJob",
             )
@@ -185,7 +185,8 @@ def assign_unassigned(cls, triggerer_id, capacity, session=None):
         # up to `capacity` of those to us.
         trigger_ids_query = (
             session.query(cls.id)
-            .filter(or_(cls.triggerer_id.notin_(alive_triggerer_ids), cls.triggerer_id == None))  # noqa: E711
+            # notin_ doesn't find NULL rows
+            .filter(or_(cls.triggerer_id.notin_(alive_triggerer_ids), cls.triggerer_id.is_(None)))
             .limit(capacity)
             .all()
         )

From e1854e35fcf59ced6dd5eed37bf3eef94bf3e87b Mon Sep 17 00:00:00 2001
From: Jake Kramer <jkramer@ginkgobioworks.com>
Date: Wed, 23 Feb 2022 16:49:17 -0500
Subject: [PATCH 3/3] Change order of `or_` params

---
 airflow/models/trigger.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/airflow/models/trigger.py b/airflow/models/trigger.py
index f394cc55b487e..28eb133162e0e 100644
--- a/airflow/models/trigger.py
+++ b/airflow/models/trigger.py
@@ -186,7 +186,7 @@ def assign_unassigned(cls, triggerer_id, capacity, session=None):
         trigger_ids_query = (
             session.query(cls.id)
             # notin_ doesn't find NULL rows
-            .filter(or_(cls.triggerer_id.notin_(alive_triggerer_ids), cls.triggerer_id.is_(None)))
+            .filter(or_(cls.triggerer_id.is_(None), cls.triggerer_id.notin_(alive_triggerer_ids)))
             .limit(capacity)
             .all()
         )
