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

1 files changed, 19 insertions, 4 deletions

diff --git a/lib/sqlalchemy/orm/loading.py b/lib/sqlalchemy/orm/loading.py
index 3599aa3e7..1728b2d37 100644
--- a/lib/sqlalchemy/orm/loading.py
+++ b/lib/sqlalchemy/orm/loading.py
@@ -180,23 +180,37 @@ def load_on_ident(query, key,
     else:
         ident = None
 
+    return load_on_pk_identity(
+        query, ident, refresh_state=refresh_state,
+        with_for_update=with_for_update,
+        only_load_props=only_load_props
+    )
+
+
+def load_on_pk_identity(query, primary_key_identity,
+                        refresh_state=None, with_for_update=None,
+                        only_load_props=None):
+
+    """Load the given primary key identity from the database."""
+
     if refresh_state is None:
         q = query._clone()
         q._get_condition()
     else:
         q = query._clone()
 
-    if ident is not None:
+    if primary_key_identity is not None:
         mapper = query._mapper_zero()
 
         (_get_clause, _get_params) = mapper._get_clause
 
         # 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
                         ])
             _get_clause = sql_util.adapt_criterion_to_null(
                 _get_clause, nones)
@@ -206,7 +220,8 @@ def load_on_ident(query, key,
 
         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)
         ])
 
         q._params = params