Add QueryEvents before_compile_update / before_compile_delete

Added new event hooks :meth:`.QueryEvents.before_compile_update` and
:meth:`.QueryEvents.before_compile_delete` which complement
:meth:`.QueryEvents.before_compile` in the case of the :meth:`.Query.update`
and :meth:`.Query.delete` methods.

Fixes: #4461
Change-Id: I47884f0e1f07d7e62870c2a918b15f917f9245ab

2 files changed, 95 insertions, 0 deletions

diff --git a/lib/sqlalchemy/orm/events.py b/lib/sqlalchemy/orm/events.py
index 5e32c6ac9..3ee0cd1f8 100644
--- a/lib/sqlalchemy/orm/events.py
+++ b/lib/sqlalchemy/orm/events.py
@@ -2335,6 +2335,82 @@ class QueryEvents(event.Events):
         The event should normally be listened with the ``retval=True``
         parameter set, so that the modified query may be returned.
 
+        .. seealso::
+
+            :meth:`.QueryEvents.before_compile_update`
+
+            :meth:`.QueryEvents.before_compile_delete`
+
+
+        """
+
+    def before_compile_update(self, query, update_context):
+        """Allow modifications to the :class:`.Query` object within
+        :meth:`.Query.update`.
+
+        Like the :meth:`.QueryEvents.before_compile` event, this event
+        should be configured with ``retval=True``, and the modified
+        :class:`.Query` object returned, as in ::
+
+            @event.listens_for(Query, "before_compile_update", retval=True)
+            def no_deleted(query, update_context):
+                for desc in query.column_descriptions:
+                    if desc['type'] is User:
+                        entity = desc['entity']
+                        query = query.filter(entity.deleted == False)
+                return query
+
+        :param query: a :class:`.Query` instance; this is also
+         the ``.query`` attribute of the given "update context"
+         object.
+
+        :param update_context: an "update context" object which is
+         the same kind of object as described in
+         :paramref:`.QueryEvents.after_bulk_update.update_context`.
+
+        .. versionadded:: 1.2.17
+
+        .. seealso::
+
+            :meth:`.QueryEvents.before_compile`
+
+            :meth:`.QueryEvents.before_compile_delete`
+
+
+        """
+
+    def before_compile_delete(self, query, delete_context):
+        """Allow modifications to the :class:`.Query` object within
+        :meth:`.Query.delete`.
+
+        Like the :meth:`.QueryEvents.before_compile` event, this event
+        should be configured with ``retval=True``, and the modified
+        :class:`.Query` object returned, as in ::
+
+            @event.listens_for(Query, "before_compile_delete", retval=True)
+            def no_deleted(query, delete_context):
+                for desc in query.column_descriptions:
+                    if desc['type'] is User:
+                        entity = desc['entity']
+                        query = query.filter(entity.deleted == False)
+                return query
+
+        :param query: a :class:`.Query` instance; this is also
+         the ``.query`` attribute of the given "delete context"
+         object.
+
+        :param delete_context: a "delete context" object which is
+         the same kind of object as described in
+         :paramref:`.QueryEvents.after_bulk_delete.delete_context`.
+
+        .. versionadded:: 1.2.17
+
+        .. seealso::
+
+            :meth:`.QueryEvents.before_compile`
+
+            :meth:`.QueryEvents.before_compile_update`
+
 
         """
 
diff --git a/lib/sqlalchemy/orm/persistence.py b/lib/sqlalchemy/orm/persistence.py
index 6e127f41d..1e3a8a44e 100644
--- a/lib/sqlalchemy/orm/persistence.py
+++ b/lib/sqlalchemy/orm/persistence.py
@@ -1638,6 +1638,7 @@ class BulkUD(object):
             return klass(*arg)
 
     def exec_(self):
+        self._do_before_compile()
         self._do_pre()
         self._do_pre_synchronize()
         self._do_exec()
@@ -1648,9 +1649,13 @@ class BulkUD(object):
         self.result = self.query._execute_crud(stmt, self.mapper)
         self.rowcount = self.result.rowcount
 
+    def _do_before_compile(self):
+        raise NotImplementedError()
+
     @util.dependencies("sqlalchemy.orm.query")
     def _do_pre(self, querylib):
         query = self.query
+
         self.context = querylib.QueryContext(query)
 
         if isinstance(query._entities[0], querylib._ColumnEntity):
@@ -1766,6 +1771,13 @@ class BulkUpdate(BulkUD):
             update_kwargs,
         )
 
+    def _do_before_compile(self):
+        if self.query.dispatch.before_compile_update:
+            for fn in self.query.dispatch.before_compile_update:
+                new_query = fn(self.query, self)
+                if new_query is not None:
+                    self.query = new_query
+
     @property
     def _resolved_values(self):
         values = []
@@ -1847,6 +1859,13 @@ class BulkDelete(BulkUD):
             query,
         )
 
+    def _do_before_compile(self):
+        if self.query.dispatch.before_compile_delete:
+            for fn in self.query.dispatch.before_compile_delete:
+                new_query = fn(self.query, self)
+                if new_query is not None:
+                    self.query = new_query
+
     def _do_exec(self):
         delete_stmt = sql.delete(self.primary_table, self.context.whereclause)