Merge "Don't cache a query that has before_compile modifications"

3 files changed, 50 insertions, 9 deletions

diff --git a/lib/sqlalchemy/ext/baked.py b/lib/sqlalchemy/ext/baked.py
index 44e28d045..d18a35a40 100644
--- a/lib/sqlalchemy/ext/baked.py
+++ b/lib/sqlalchemy/ext/baked.py
@@ -225,6 +225,7 @@ class BakedQuery(object):
         query = self._as_query(session)
 
         context = query._compile_context()
+
         self._bake_subquery_loaders(session, context)
         context.session = None
         context.query = query = context.query.with_session(None)
@@ -242,7 +243,13 @@ class BakedQuery(object):
             "_joinpoint",
         ):
             query.__dict__.pop(attr, None)
-        self._bakery[self._effective_key(session)] = context
+
+        # if the query is not safe to cache, we still do everything as though
+        # we did cache it, since the receiver of _bake() assumes subqueryload
+        # context was set up, etc.
+        if context.query._bake_ok:
+            self._bakery[self._effective_key(session)] = context
+
         return context
 
     def to_query(self, query_or_session):
@@ -332,6 +339,9 @@ class BakedQuery(object):
         like a Query object.
 
         """
+        if "baked_queries" not in context.attributes:
+            return
+
         for k, cache_key, query in context.attributes["baked_queries"]:
             bk = BakedQuery(
                 self._bakery, lambda sess, q=query: q.with_session(sess)
diff --git a/lib/sqlalchemy/orm/events.py b/lib/sqlalchemy/orm/events.py
index 5bb67b68f..2998a7639 100644
--- a/lib/sqlalchemy/orm/events.py
+++ b/lib/sqlalchemy/orm/events.py
@@ -2397,12 +2397,31 @@ class QueryEvents(event.Events):
         The event should normally be listened with the ``retval=True``
         parameter set, so that the modified query may be returned.
 
-        .. warning::  If the :meth:`.QueryEvents.before_compile` event is to
-           be applied to :class:`.Query` objects that are used for lazy loading
-           of :func:`.relationships` (as described at :ref:`lazy_loading`),
-           it may be necessary to set :paramref:`.relationship.bake_queries`
-           to ``False``, else the :meth:`.QueryEvents.before_compile` event
-           will not be invoked for each lazy load operation.
+        The :meth:`.QueryEvents.before_compile` event by default
+        will disallow "baked" queries from caching a query, if the event
+        hook returns a new :class:`.Query` object.   This affects both direct
+        use of the baked query extension as well as its operation within
+        lazy loaders and eager loaders for relationships.  In order to
+        re-establish the query being cached, apply the event adding the
+        ``bake_ok`` flag::
+
+            @event.listens_for(
+                Query, "before_compile", retval=True, bake_ok=True)
+            def my_event(query):
+                for desc in query.column_descriptions:
+                    if desc['type'] is User:
+                        entity = desc['entity']
+                        query = query.filter(entity.deleted == False)
+                return query
+
+        When ``bake_ok`` is set to True, the event hook will only be invoked
+        once, and not called for subsequent invocations of a particular query
+        that is being cached.
+
+        .. versionadded:: 1.3.11  - added the "bake_ok" flag to the
+           :meth:`.QueryEvents.before_compile` event and disallowed caching via
+           the "baked" extension from occurring for event handlers that
+           return  a new :class:`.Query` object if this flag is not set.
 
         .. seealso::
 
@@ -2410,6 +2429,7 @@ class QueryEvents(event.Events):
 
             :meth:`.QueryEvents.before_compile_delete`
 
+            :ref:`baked_with_before_compile`
 
         """
 
@@ -2494,7 +2514,7 @@ class QueryEvents(event.Events):
         """
 
     @classmethod
-    def _listen(cls, event_key, retval=False, **kw):
+    def _listen(cls, event_key, retval=False, bake_ok=False, **kw):
         fn = event_key._listen_fn
 
         if not retval:
@@ -2508,5 +2528,13 @@ class QueryEvents(event.Events):
                     return fn(*arg, **kw)
 
             event_key = event_key.with_wrapper(wrap)
+        else:
+            # don't assume we can apply an attribute to the callable
+            def wrap(*arg, **kw):
+                return fn(*arg, **kw)
+
+            event_key = event_key.with_wrapper(wrap)
+
+        wrap._bake_ok = bake_ok
 
         event_key.base_listen(**kw)
diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py
index 35416ae5e..45042ed33 100644
--- a/lib/sqlalchemy/orm/query.py
+++ b/lib/sqlalchemy/orm/query.py
@@ -127,6 +127,7 @@ class Query(Generative):
     _orm_only_from_obj_alias = True
     _current_path = _path_registry
     _has_mapper_entities = False
+    _bake_ok = True
 
     lazy_loaded_from = None
     """An :class:`.InstanceState` that is using this :class:`.Query` for a
@@ -3812,8 +3813,10 @@ class Query(Generative):
         if self.dispatch.before_compile:
             for fn in self.dispatch.before_compile:
                 new_query = fn(self)
-                if new_query is not None:
+                if new_query is not None and new_query is not self:
                     self = new_query
+                    if not fn._bake_ok:
+                        self._bake_ok = False
 
         context = QueryContext(self)