<s> """<STR_LIT>""" <EOL> from . . import util <EOL> from . . orm . session import Session <EOL> from . . orm . query import Query <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> class ShardedQuery ( Query ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> super ( ShardedQuery , self ) . __init__ ( * args , ** kwargs ) <EOL> self . id_chooser = self . session . id_chooser <EOL> self . query_chooser = self . session . query_chooser <EOL> self . _shard_id = None <EOL> def set_shard ( self , shard_id ) : <EOL> """<STR_LIT>""" <EOL> q = self . _clone ( ) <EOL> q . _shard_id = shard_id <EOL> return q <EOL> def _execute_and_instances ( self , context ) : <EOL> def iter_for_shard ( shard_id ) : <EOL> context . attributes [ '<STR_LIT>' ] = shard_id <EOL> result = self . _connection_from_session ( <EOL> mapper = self . _mapper_zero ( ) , <EOL> shard_id = shard_id ) . execute ( <EOL> context . statement , <EOL> self . _params ) <EOL> return self . instances ( result , context ) <EOL> if self . _shard_id is not None : <EOL> return iter_for_shard ( self . _shard_id ) <EOL> else : <EOL> partial = [ ] <EOL> for shard_id in self . query_chooser ( self ) : <EOL> partial . extend ( iter_for_shard ( shard_id ) ) <EOL> return iter ( partial ) <EOL> def get ( self , ident , ** kwargs ) : <EOL> if self . _shard_id is not None : <EOL> return super ( ShardedQuery , self ) . get ( ident ) <EOL> else : <EOL> ident = util . to_list ( ident ) <EOL> for shard_id in self . id_chooser ( self , ident ) : <EOL> o = self . set_shard ( shard_id ) . get ( ident , ** kwargs ) <EOL> if o is not None : <EOL> return o <EOL> else : <EOL> return None <EOL> class ShardedSession ( Session ) : <EOL> def __init__ ( self , shard_chooser , id_chooser , query_chooser , shards = None , <EOL> query_cls = ShardedQuery , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> super ( ShardedSession , self ) . __init__ ( query_cls = query_cls , ** kwargs ) <EOL> self . shard_chooser = shard_chooser <EOL> self . id_chooser = id_chooser <EOL> self . query_chooser = query_chooser <EOL> self . __binds = { } <EOL> self . connection_callable = self . connection <EOL> if shards is not None : <EOL> for k in shards : <EOL> self . bind_shard ( k , shards [ k ] ) <EOL> def connection ( self , mapper = None , instance = None , shard_id = None , ** kwargs ) : <EOL> if shard_id is None : <EOL> shard_id = self . shard_chooser ( mapper , instance ) <EOL> if self . transaction is not None : <EOL> return self . transaction . connection ( mapper , shard_id = shard_id ) <EOL> else : <EOL> return self . get_bind ( <EOL> mapper , <EOL> shard_id = shard_id , <EOL> instance = instance <EOL> ) . contextual_connect ( ** kwargs ) <EOL> def get_bind ( self , mapper , shard_id = None , <EOL> instance = None , clause = None , ** kw ) : <EOL> if shard_id is None : <EOL> shard_id = self . shard_chooser ( mapper , instance , clause = clause ) <EOL> return self . __binds [ shard_id ] <EOL> def bind_shard ( self , shard_id , bind ) : <EOL> self . __binds [ shard_id ] = bind </s>
