kaxil commented on code in PR #70961:
URL: https://github.com/apache/airflow/pull/70961#discussion_r4042435956


##########
airflow-core/src/airflow/models/trigger.py:
##########
@@ -231,42 +231,58 @@ def fetch_trigger_ids_with_non_task_associations(cls, *, 
session: Session = NEW_
         return set(session.scalars(query))
 
     @classmethod
-    @provide_session
-    def clean_unused(cls, *, session: Session = NEW_SESSION) -> None:
+    def clean_unused(cls) -> None:
         """
-        Delete all triggers that have no tasks dependent on them and are not 
associated to an asset.
+        Delete triggers that have no dependent tasks, assets, or callbacks in 
bounded transactions.
 
         Triggers have a one-to-many relationship to task instances, so we need 
to clean those up first.
         Afterward we can drop the triggers not referenced by anyone.
         """
-        # Update all task instances with trigger IDs that are not DEFERRED to 
remove them
-        for attempt in run_with_db_retries():
-            with attempt:
-                session.execute(
-                    update(TaskInstance)
-                    .where(
-                        TaskInstance.state != TaskInstanceState.DEFERRED, 
TaskInstance.trigger_id.is_not(None)
-                    )
-                    .values(trigger_id=None)
-                )
-
-        # Get all triggers that have no task instances, assets, or callbacks 
depending on them and delete them
-        ids = select(cls.id).where(
-            ~cls.assets.any(),
-            ~cls.callback.has(),
-            ~cls.task_instance.has(),
+        batch_size = conf.getint("triggerer", 
"unreferenced_triggers_cleanup_batch_size", fallback=500)
+        if batch_size <= 0:
+            raise ValueError("[triggerer] 
unreferenced_triggers_cleanup_batch_size must be at least 1")
+
+        clear_task_instance_references = True
+        while True:
+            deleted_count = 0
+            for attempt in run_with_db_retries():
+                with attempt:
+                    with create_session(scoped=False) as session:
+                        if clear_task_instance_references:
+                            # Update all task instances with trigger IDs that 
are not DEFERRED to remove them
+                            session.execute(
+                                update(TaskInstance)
+                                .where(
+                                    TaskInstance.state != 
TaskInstanceState.DEFERRED,
+                                    TaskInstance.trigger_id.is_not(None),
+                                )
+                                .values(trigger_id=None)
+                            )
+                        deleted_count = cls._delete_unused_batch(batch_size, 
session=session)
+
+            clear_task_instance_references = False
+            if deleted_count < batch_size:
+                return
+
+    @classmethod
+    def _delete_unused_batch(cls, batch_size: int, *, session: Session) -> int:
+        ids = (
+            select(cls.id)
+            .where(
+                ~cls.assets.any(),
+                ~cls.callback.has(),
+                ~cls.task_instance.has(),
+            )
+            .order_by(cls.id)
+            .limit(batch_size)
         )
         ids = with_row_locks(ids, session, of=cls, skip_locked=True, 
key_share=False)
-        if get_dialect_name(session) == "mysql":
-            # MySQL doesn't support DELETE with JOIN, so we need to do it in 
two steps
-            ids_list = list(session.scalars(ids).all())
-            session.execute(
-                
delete(Trigger).where(Trigger.id.in_(ids_list)).execution_options(synchronize_session=False)
-            )
-        else:
-            session.execute(
-                
delete(Trigger).where(Trigger.id.in_(ids)).execution_options(synchronize_session=False)
-            )
+        ids_list = list(session.scalars(ids))
+        if not ids_list:
+            return 0
+
+        
session.execute(delete(cls).where(cls.id.in_(ids_list)).execution_options(synchronize_session=False))

Review Comment:
   Worth a look before this lands: dropping the dialect branch also drops the 
one thing that made the non-MySQL path safe.
   
   Previously the `else` branch passed the `ids` select straight into the 
DELETE, so the three `~...any()/has()` predicates were re-evaluated inside the 
statement. This now materialises `ids_list` for every backend and deletes by id 
alone, which is the MySQL-only shape that #72061 reports as deleting task 
instances.
   
   The window is between the SELECT and this DELETE. `with_row_locks` covers it 
while `[scheduler] use_row_level_locking` is on, but it returns the query 
unchanged when that is off (`airflow/utils/sqlalchemy.py`), and then nothing 
holds those rows. A referrer attaching in the gap is reachable: the Dag 
processor reuses committed trigger rows by `(classpath, kwargs)` when it 
registers asset watchers (`dag_processing/collection.py`). Both 
`task_instance.trigger_id` and `asset_watcher.trigger_id` are `ON DELETE 
CASCADE`, so deleting a re-referenced trigger takes the referring row with it.
   
   Cheap fix that keeps the batching, re-applying the same predicates in the 
DELETE:
   
   ```python
   unreferenced = (~cls.assets.any(), ~cls.callback.has(), 
~cls.task_instance.has())
   ids = select(cls.id).where(*unreferenced).order_by(cls.id).limit(batch_size)
   ...
   session.execute(delete(cls).where(cls.id.in_(ids_list), 
*unreferenced).execution_options(synchronize_session=False))
   ```
   
   These correlate to `task_instance`, `callback` and `asset_watcher` rather 
than to `trigger`, so it does not hit the MySQL restriction that motivated the 
two-step in the first place. I checked that it compiles and runs on MySQL 8.0.
   



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to