Don't cache a query that has before_compile modifications

The :class:`.BakedQuery` will not cache a query that was modified by a
:meth:`.QueryEvents.before_compile` event, so that compilation hooks that
may be applying ad-hoc modifications to queries will take effect on each
run.  In particular this is helpful for events that modify queries used in
lazy loading as well as eager loading such as "select in" loading.  In
order to re-enable caching for a query modified by this event, a new
flag ``bake_ok`` is added; see :ref:`baked_with_before_compile` for
details.

A longer term plan to provide a new form of SQL caching should solve this
kind of issue more comprehensively.

Fixes: #4947
Change-Id: I5823c4fa00e7b6d46a2e8461b02d8b16605a6ed0

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 92f9ee952..fc443d053 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)