- Horizontal shard query places 'shard_id' in

context.attributes where it's accessible by the
"load()" event. [ticket:2031]

1 files changed, 2 insertions, 0 deletions

diff --git a/lib/sqlalchemy/ext/horizontal_shard.py b/lib/sqlalchemy/ext/horizontal_shard.py
index 41fae8e7b..926c029d0 100644
--- a/lib/sqlalchemy/ext/horizontal_shard.py
+++ b/lib/sqlalchemy/ext/horizontal_shard.py
@@ -95,6 +95,7 @@ class ShardedQuery(Query):
 
     def _execute_and_instances(self, context):
         if self._shard_id is not None:
+            context.attributes['shard_id'] = self._shard_id
             result = self.session.connection(
                             mapper=self._mapper_zero(),
                             shard_id=self._shard_id).execute(context.statement, self._params)
@@ -102,6 +103,7 @@ class ShardedQuery(Query):
         else:
             partial = []
             for shard_id in self.query_chooser(self):
+                context.attributes['shard_id'] = shard_id
                 result = self.session.connection(
                             mapper=self._mapper_zero(),
                             shard_id=shard_id).execute(context.statement, self._params)