From: Mike Bayer Date: Wed, 23 Jan 2019 02:49:07 +0000 (-0500) Subject: Add QueryEvents before_compile_update / before_compile_delete X-Git-Tag: rel_1_3_0b2~15 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=bd735eba637cbf2c157046f72dc795a8b2b803e7;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git 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 --- diff --git a/doc/build/changelog/unreleased_12/4461.rst b/doc/build/changelog/unreleased_12/4461.rst new file mode 100644 index 0000000000..6ca4ab8ae2 --- /dev/null +++ b/doc/build/changelog/unreleased_12/4461.rst @@ -0,0 +1,9 @@ +.. change:: + :tags: feature, orm + :tickets: 4461 + + 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. + diff --git a/lib/sqlalchemy/orm/events.py b/lib/sqlalchemy/orm/events.py index 5e32c6ac9f..3ee0cd1f85 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 6e127f41d8..1e3a8a44ea 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) diff --git a/test/orm/test_events.py b/test/orm/test_events.py index bb1a935de3..bf2c6c1493 100644 --- a/test/orm/test_events.py +++ b/test/orm/test_events.py @@ -25,6 +25,7 @@ from sqlalchemy.testing import AssertsCompiledSQL from sqlalchemy.testing import eq_ from sqlalchemy.testing import fixtures from sqlalchemy.testing import is_not_ +from sqlalchemy.testing.assertsql import CompiledSQL from sqlalchemy.testing.mock import ANY from sqlalchemy.testing.mock import call from sqlalchemy.testing.mock import Mock @@ -2852,7 +2853,10 @@ class SessionExtensionTest(_fixtures.FixtureTest): class QueryEventsTest( - _RemoveListeners, _fixtures.FixtureTest, AssertsCompiledSQL + _RemoveListeners, + _fixtures.FixtureTest, + AssertsCompiledSQL, + testing.AssertsExecutionResults, ): __dialect__ = "default" @@ -2903,6 +2907,55 @@ class QueryEventsTest( ) eq_(q.all(), [(7, "jack")]) + def test_before_compile_update(self): + @event.listens_for(query.Query, "before_compile_update", retval=True) + def no_deleted(query, update_context): + assert update_context.query is query + + for desc in query.column_descriptions: + if desc["type"] is User: + entity = desc["expr"] + query = query.filter(entity.id != 10) + return query + + User = self.classes.User + s = Session() + + with self.sql_execution_asserter() as asserter: + q = s.query(User).filter_by(id=7).update({"name": "ed"}) + asserter.assert_( + CompiledSQL( + "UPDATE users SET name=:name WHERE " + "users.id = :id_1 AND users.id != :id_2", + [{"name": "ed", "id_1": 7, "id_2": 10}], + ) + ) + + def test_before_compile_delete(self): + @event.listens_for(query.Query, "before_compile_delete", retval=True) + def no_deleted(query, delete_context): + assert delete_context.query is query + + for desc in query.column_descriptions: + if desc["type"] is User: + entity = desc["expr"] + query = query.filter(entity.id != 10) + return query + + User = self.classes.User + s = Session() + + # note this deletes no rows + with self.sql_execution_asserter() as asserter: + q = s.query(User).filter_by(id=10).delete() + asserter.assert_( + CompiledSQL( + "DELETE FROM users WHERE " + "users.id = :id_1 AND users.id != :id_2", + [{"id_1": 10, "id_2": 10}], + ) + ) + class RefreshFlushInReturningTest(fixtures.MappedTest): """test [ticket:3427].