diff options
Diffstat (limited to 'lib/sqlalchemy')
| -rw-r--r-- | lib/sqlalchemy/ext/horizontal_shard.py | 19 | ||||
| -rw-r--r-- | lib/sqlalchemy/orm/query.py | 19 |
2 files changed, 26 insertions, 12 deletions
diff --git a/lib/sqlalchemy/ext/horizontal_shard.py b/lib/sqlalchemy/ext/horizontal_shard.py index dfd471c78..6aafb2274 100644 --- a/lib/sqlalchemy/ext/horizontal_shard.py +++ b/lib/sqlalchemy/ext/horizontal_shard.py @@ -40,20 +40,21 @@ class ShardedQuery(Query): return q def _execute_and_instances(self, context): - if self._shard_id is not None: - context.attributes['shard_id'] = self._shard_id - result = self.session.connection( + def iter_for_shard(shard_id): + context.attributes['shard_id'] = shard_id + result = self._connection_from_session( mapper=self._mapper_zero(), - shard_id=self._shard_id).execute(context.statement, self._params) + shard_id=shard_id).execute( + context.statement, + self._params) return self.instances(result, context) + + if self._shard_id is not None: + return iter_for_shard(self._shard_id) 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) - partial = partial + list(self.instances(result, context)) + partial.extend(iter_for_shard(shard_id)) # if some kind of in memory 'sorting' # were done, this is where it would happen diff --git a/lib/sqlalchemy/orm/query.py b/lib/sqlalchemy/orm/query.py index ef42e0d3a..75fd5870e 100644 --- a/lib/sqlalchemy/orm/query.py +++ b/lib/sqlalchemy/orm/query.py @@ -773,6 +773,14 @@ class Query(object): m = _MapperEntity(self, entity) self._setup_aliasizers([m]) + @_generative() + def with_session(self, session): + """Return a :class:`Query` that will use the given :class:`.Session`. + + """ + + self.session = session + def from_self(self, *entities): """return a Query that selects from this Query's SELECT statement. @@ -1766,13 +1774,18 @@ class Query(object): self.session._autoflush() return self._execute_and_instances(context) - def _execute_and_instances(self, querycontext): + def _connection_from_session(self, **kw): conn = self.session.connection( + **kw) + if self._execution_options: + conn = conn.execution_options(**self._execution_options) + return conn + + def _execute_and_instances(self, querycontext): + conn = self._connection_from_session( mapper = self._mapper_zero_or_none(), clause = querycontext.statement, close_with_result=True) - if self._execution_options: - conn = conn.execution_options(**self._execution_options) result = conn.execute(querycontext.statement, self._params) return self.instances(result, querycontext) |