#725 add query arg to id_chooser()

diff --git a/examples/sharding/attribute_shard.py b/examples/sharding/attribute_shard.py
index 6e4732989646229b4c331cb4f98e6b36a2b325f7..df3f7467f736c30367c84b32bd5b9baa69abffc9 100644 (file)
--- a/examples/sharding/attribute_shard.py
+++ b/examples/sharding/attribute_shard.py
@@ -118,7 +118,7 @@ def shard_chooser(mapper, instance):
 # pk so we just return all shard ids. often, youd want to do some 
 # kind of round-robin strategy here so that requests are evenly 
 # distributed among DBs
-def id_chooser(ident):
+def id_chooser(query, ident):
     return ['north_america', 'asia', 'europe', 'south_america']
 
 # query_chooser.  this also returns a list of shard ids, which can

diff --git a/lib/sqlalchemy/orm/shard.py b/lib/sqlalchemy/orm/shard.py
index 69d692167362d29455acbc7022774678b9f7edda..7f4f3bd210e190116dbc44c998e91d4f46018c43 100644 (file)
--- a/lib/sqlalchemy/orm/shard.py
+++ b/lib/sqlalchemy/orm/shard.py
@@ -15,9 +15,10 @@ class ShardedSession(Session):
                 the future as participating in that shard.
             
             id_chooser
-                a callable, passed a tuple of identity values, which should return
-                a list of shard ids where the ID might reside.  The databases will
-                be queried in the order of this listing.
+                a callable, passed a query and a tuple of identity values,
+                which should return a list of shard ids where the ID might
+                reside.  The databases will be queried in the order of this
+                listing.
                 
             query_chooser
                 for a given Query, returns the list of shard_ids where the query
@@ -98,7 +99,7 @@ class ShardedQuery(Query):
         if self._shard_id is not None:
             return super(ShardedQuery, self).get(ident)
         else:
-            for shard_id in self.id_chooser(ident):
+            for shard_id in self.id_chooser(self, ident):
                 o = self.set_shard(shard_id).get(ident, **kwargs)
                 if o is not None:
                     return o
@@ -109,7 +110,7 @@ class ShardedQuery(Query):
         if self._shard_id is not None:
             return super(ShardedQuery, self).load(ident)
         else:
-            for shard_id in self.id_chooser(ident):
+            for shard_id in self.id_chooser(self, ident):
                 o = self.set_shard(shard_id).load(ident, raiseerr=False, **kwargs)
                 if o is not None:
                     return o

diff --git a/test/orm/sharding/shard.py b/test/orm/sharding/shard.py
index c1dd63d6575e72aa367d74baa8c9a41ad9e16818..69839c8c7ad742691c4ac7fb714d4c586e5bfa4e 100644 (file)
--- a/test/orm/sharding/shard.py
+++ b/test/orm/sharding/shard.py
@@ -70,7 +70,7 @@ class ShardTest(PersistTest):
             else:
                 return shard_chooser(mapper, instance.location)
 
-        def id_chooser(ident):
+        def id_chooser(query, ident):
             return ['north_america', 'asia', 'europe', 'south_america']
 
         def query_chooser(query):