Add Query.lazy_load_from attribute for sharding

Added new attribute :attr:`.Query.lazy_loaded_from` which is populated
with an :class:`.InstanceState` that is using this :class:`.Query` in
order to lazy load a relationship.  The rationale for this is that
it serves as a hint for the horizontal sharding feature to use, such that
the identity token of the state can be used as the default identity token
to use for the query within id_chooser().

Also repaired an issue in the :meth:`.Result.with_post_criteria`
method added in I899808734458e25a023142c2c5bb37cbed869479
for :ticket:`4128` where the "unbake subquery loaders" version was calling
the post crtieria functions given the :class:`.Result` as the argument
rather than applying them to the :class:`.Query`.

Change-Id: I3c0919ce7fd151b80fe2f9b5f99f60df31c2d73d
Fixes: #4243

1 files changed, 7 insertions, 1 deletions

diff --git a/examples/sharding/attribute_shard.py b/examples/sharding/attribute_shard.py
index cd9b14d5e..0e19b69f3 100644
--- a/examples/sharding/attribute_shard.py
+++ b/examples/sharding/attribute_shard.py
@@ -133,7 +133,13 @@ def id_chooser(query, ident):
     distributed among DBs.
 
     """
-    return ['north_america', 'asia', 'europe', 'south_america']
+    if query.lazy_loaded_from:
+        # if we are in a lazy load, we can look at the parent object
+        # and limit our search to that same shard, assuming that's how we've
+        # set things up.
+        return [query.lazy_loaded_from.identity_token]
+    else:
+        return ['north_america', 'asia', 'europe', 'south_america']
 
 
 def query_chooser(query):