Refactor "get" to allow for pluggable identity token schemes

Fixed regression in 1.2 within sharded query feature where the
new "identity_token" element was not being correctly considered within
the scope of a lazy load operation, when searching the identity map
for a related many-to-one element.   The new behavior will allow for
making use of the "id_chooser" in order to determine the best identity
key to retrieve from the identity map.  In order to achieve this, some
refactoring of 1.2's "identity_token" approach has made some slight changes
to the implementation of ``ShardedQuery`` which should be noted for other
derivations of this class.

Change-Id: I04fa60535deec2d0cdec89f602935dfebeb9eb9d
Fixes: #4228

2 files changed, 57 insertions, 21 deletions

diff --git a/lib/sqlalchemy/ext/baked.py b/lib/sqlalchemy/ext/baked.py
index 86eee831b..f4d71f410 100644
--- a/lib/sqlalchemy/ext/baked.py
+++ b/lib/sqlalchemy/ext/baked.py
@@ -446,12 +446,10 @@ class Result(object):
         """
 
         query = self.bq.steps[0](self.session)
-        return query._get_impl(ident, self._load_on_ident)
+        return query._get_impl(ident, self._load_on_pk_identity)
 
-    def _load_on_ident(self, query, key):
-        """Load the given identity key from the database."""
-
-        ident = key[1]
+    def _load_on_pk_identity(self, query, primary_key_identity):
+        """Load the given primary key identity from the database."""
 
         mapper = query._mapper_zero()
 
@@ -465,10 +463,11 @@ class Result(object):
 
             # None present in ident - turn those comparisons
             # into "IS NULL"
-            if None in ident:
+            if None in primary_key_identity:
                 nones = set([
                     _get_params[col].key for col, value in
-                    zip(mapper.primary_key, ident) if value is None
+                    zip(mapper.primary_key, primary_key_identity)
+                    if value is None
                 ])
                 _lcl_get_clause = sql_util.adapt_criterion_to_null(
                     _lcl_get_clause, nones)
@@ -490,11 +489,13 @@ class Result(object):
         bq = bq._clone()
         bq._cache_key += (_get_clause, )
 
-        bq = bq.with_criteria(setup, tuple(elem is None for elem in ident))
+        bq = bq.with_criteria(
+            setup, tuple(elem is None for elem in primary_key_identity))
 
         params = dict([
             (_get_params[primary_key].key, id_val)
-            for id_val, primary_key in zip(ident, mapper.primary_key)
+            for id_val, primary_key
+            in zip(primary_key_identity, mapper.primary_key)
         ])
 
         result = list(bq.for_session(self.session).params(**params))
diff --git a/lib/sqlalchemy/ext/horizontal_shard.py b/lib/sqlalchemy/ext/horizontal_shard.py
index 039a4656b..266bd784e 100644
--- a/lib/sqlalchemy/ext/horizontal_shard.py
+++ b/lib/sqlalchemy/ext/horizontal_shard.py
@@ -62,30 +62,65 @@ class ShardedQuery(Query):
             # were done, this is where it would happen
             return iter(partial)
 
-    def _get_impl(self, ident, fallback_fn):
-        # TODO: the "ident" here should be getting the identity token
-        # which indicates that this area can likely be simplified, as the
-        # token will fall through into _execute_and_instances
-        def _fallback(query, ident):
+    @classmethod
+    def _identity_lookup(
+            cls, session, mapper, primary_key_identity, identity_token=None,
+            **kw):
+        """override the default Query._identity_lookup method so that we
+        search for a given non-token primary key identity across all
+        possible identity tokens (e.g. shard ids).
+
+        """
+
+        if identity_token is not None:
+            return super(ShardedQuery, cls)._identity_lookup(
+                session, mapper, primary_key_identity,
+                identity_token=identity_token,
+                **kw
+            )
+        else:
+            q = cls([mapper], session)
+            for shard_id in q.id_chooser(q, primary_key_identity):
+                obj = super(ShardedQuery, cls)._identity_lookup(
+                    session, mapper, primary_key_identity,
+                    identity_token=shard_id,
+                    **kw
+                )
+                if obj is not None:
+                    return obj
+
+            return None
+
+    def _get_impl(
+            self, primary_key_identity, db_load_fn, identity_token=None):
+        """Override the default Query._get_impl() method so that we emit
+        a query to the DB for each possible identity token, if we don't
+        have one already.
+
+        """
+        def _db_load_fn(query, primary_key_identity):
+            # load from the database.  The original db_load_fn will
+            # use the given Query object to load from the DB, so our
+            # shard_id is what will indicate the DB that we query from.
             if self._shard_id is not None:
-                return fallback_fn(self, ident)
+                return db_load_fn(self, primary_key_identity)
             else:
-                ident = util.to_list(ident)
+                ident = util.to_list(primary_key_identity)
+                # build a ShardedQuery for each shard identifier and
+                # try to load from the DB
                 for shard_id in self.id_chooser(self, ident):
                     q = self.set_shard(shard_id)
-                    o = fallback_fn(q, ident)
+                    o = db_load_fn(q, ident)
                     if o is not None:
                         return o
                 else:
                     return None
 
-        if self._shard_id is not None:
+        if identity_token is None and self._shard_id is not None:
             identity_token = self._shard_id
-        else:
-            identity_token = None
 
         return super(ShardedQuery, self)._get_impl(
-            ident, _fallback, identity_token=identity_token)
+            primary_key_identity, _db_load_fn, identity_token=identity_token)
 
 
 class ShardedSession(Session):