- fixes to ShardedSession to work with deferred columns [ticket:771].

- user-defined shard_chooser() function must accept "clause=None"
  argument; this is the ClauseElement passed to session.execute(statement)
  and can be used to determine correct shard id (since execute() doesn't
  take an instance)

3 files changed, 34 insertions, 19 deletions

diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index b616570ab..6f06474b7 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -488,20 +488,19 @@ class Session(object):
 
         If this ``Session`` is transactional, the connection will be in
         the context of this session's transaction.  Otherwise, the
-        connection is returned by the ``contextual_connect()`` method, which
-        some Engines override to return a thread-local connection, and
-        will have `close_with_result` set to `True`.
-
-        The given `**kwargs` will be sent to the engine's
-        ``contextual_connect()`` method, if no transaction is in
-        progress.
+        connection is returned by the ``contextual_connect()`` method
+        on the engine.
         
         the "mapper" argument is a class or mapper to which a bound engine
         will be located; use this when the Session itself is either bound
         to multiple engines or connections, or is not bound to any connectable.
+        
+        \**kwargs are additional arguments which will be passed to get_bind().
+        See the get_bind() method for details.  Note that the "ShardedSession"
+        subclass takes a different get_bind() argument signature.
         """
 
-        return self.__connection(self.get_bind(mapper))
+        return self.__connection(self.get_bind(mapper, **kwargs))
 
     def __connection(self, engine, **kwargs):
         if self.transaction is not None:
@@ -592,8 +591,23 @@ class Session(object):
 
         self.__binds[table] = bind
 
-    def get_bind(self, mapper, clause=None):
-
+    def get_bind(self, mapper, clause=None, **kwargs):
+        """return an engine corresponding to the given arguments.
+        
+            mapper
+                mapper relative to the desired operation
+            
+            clause
+                a ClauseElement which is to be executed.  if
+                mapper is not present, this may be used to locate
+                Table objects, which are then associated with mappers
+                which have associated binds.
+                
+            \**kwargs
+                Subclasses (i.e. ShardedSession) may add additional arguments 
+                to get_bind() which are passed through here.
+        """
+        
         if mapper is None and clause is None:
             if self.bind is not None:
                 return self.bind
diff --git a/lib/sqlalchemy/orm/shard.py b/lib/sqlalchemy/orm/shard.py
index 752806c0c..48e057966 100644
--- a/lib/sqlalchemy/orm/shard.py
+++ b/lib/sqlalchemy/orm/shard.py
@@ -9,11 +9,11 @@ class ShardedSession(Session):
         """construct a ShardedSession.
         
             shard_chooser
-                a callable which, passed a Mapper and a mapped instance, returns a
-                shard ID.  this id may be based off of the attributes present within the
-                object, or on some round-robin scheme.  If the scheme is based on a
-                selection, it should set whatever state on the instance to mark it in
-                the future as participating in that shard.
+                a callable which, passed a Mapper, a mapped instance, and possibly a
+                SQL clause, returns a shard ID. this id may be based off of the
+                attributes present within the object, or on some round-robin scheme. If
+                the scheme is based on a selection, it should set whatever state on the
+                instance to mark it in the future as participating in that shard.
             
             id_chooser
                 a callable, passed a query and a tuple of identity values,
@@ -47,9 +47,9 @@ class ShardedSession(Session):
         else:
             return self.get_bind(mapper, shard_id=shard_id, instance=instance).contextual_connect(**kwargs)
     
-    def get_bind(self, mapper, shard_id=None, instance=None):
+    def get_bind(self, mapper, shard_id=None, instance=None, clause=None):
         if shard_id is None:
-            shard_id = self.shard_chooser(mapper, instance)
+            shard_id = self.shard_chooser(mapper, instance, clause=clause)
         return self.__binds[shard_id]
 
     def bind_shard(self, shard_id, bind):
diff --git a/lib/sqlalchemy/orm/strategies.py b/lib/sqlalchemy/orm/strategies.py
index 86be77ae0..1ece80616 100644
--- a/lib/sqlalchemy/orm/strategies.py
+++ b/lib/sqlalchemy/orm/strategies.py
@@ -196,8 +196,9 @@ class DeferredColumnLoader(LoaderStrategy):
                 statement = sql.select([p.columns[0] for p in group], clause, from_obj=[localparent.mapped_table], use_labels=True)
             else:
                 statement, params = create_statement()
-                
-            result = session.execute(statement, params, mapper=localparent)
+            
+            conn = session.connection(mapper=localparent, instance=instance)
+            result = conn.execute(statement, params)
             try:
                 row = result.fetchone()
                 for prop in group: