def global_session_scope():
shards = {str(k): v for k, v in engine_manager.engines.items()}
session = ShardedSession(shard_chooser=shard_chooser,
id_chooser=id_chooser,
query_chooser=query_chooser,
shards=shards)
# STOPSHIP(emfree): need instrumentation and proper exception handling
# here.
try:
yield session
finally:
session.close()
def global_session_scope():
shards = {str(k): v for k, v in engine_manager.engines.items()}
session = ShardedSession(
shard_chooser=shard_chooser,
id_chooser=id_chooser,
query_chooser=query_chooser,
shards=shards)
# STOPSHIP(emfree): need instrumentation and proper exception handling
# here.
try:
yield session
finally:
session.close()
def test_query_chooser(self):
m1 = mock.Mock()
with testing.expect_deprecated(
"The ``query_choser`` parameter is deprecated; please use"):
s = ShardedSession(
shard_chooser=m1.shard_chooser,
id_chooser=m1.id_chooser,
query_chooser=m1.query_chooser,
)
m2 = mock.Mock()
s.execute_chooser(m2)
eq_(m1.mock_calls, [mock.call.query_chooser(m2.statement)])
def _session_fixture(self, **kw):
return ShardedSession(shards={"main": testing.db},
shard_chooser=lambda *args: "main",
id_chooser=lambda *args: ["fake", "main"],
query_chooser=lambda *args: ["fake", "main"],
**kw)
def _session_fixture(self):
return ShardedSession(shards={
"main": testing.db,
},
shard_chooser=lambda *args: 'main',
id_chooser=lambda *args: ['fake', 'main'],
query_chooser=lambda *args: ['fake', 'main'])
def _session_fixture(self, **kw):
# the "fake" key here is to ensure that neither id_chooser
# nor query_chooser are actually used, only shard_chooser
# should be used.
return ShardedSession(shards={"main": testing.db},
shard_chooser=lambda *args: "main",
id_chooser=lambda *args: ["fake", "main"],
execute_chooser=lambda *args: ["fake", "main"],
**kw)
def _fixture(self, lazy_load_book=False, lazy_load_pages=False):
Book, Page = self.classes("Book", "Page")
def shard_for_book(book):
if book.title == "title 1":
return "test"
elif book.title == "title 2":
return "test2"
else:
assert False
def id_chooser(query, ident):
assert query.lazy_loaded_from
if isinstance(query.lazy_loaded_from.obj(), Book):
token = shard_for_book(query.lazy_loaded_from.obj())
assert query.lazy_loaded_from.identity_token == token
return [query.lazy_loaded_from.identity_token]
def no_query_chooser(orm_context):
if (
orm_context.statement.column_descriptions[0]["type"] is Book
and lazy_load_book
):
assert isinstance(orm_context.lazy_loaded_from.obj(), Page)
elif (
orm_context.statement.column_descriptions[0]["type"] is Page
and lazy_load_pages
):
assert isinstance(orm_context.lazy_loaded_from.obj(), Book)
if orm_context.lazy_loaded_from is None:
return ["test", "test2"]
else:
return [orm_context.lazy_loaded_from.identity_token]
def shard_chooser(mapper, instance, **kw):
if isinstance(instance, Page):
return shard_for_book(instance.book)
else:
return shard_for_book(instance)
db1, db2 = self._init_dbs()
session = ShardedSession(
shards={"test": db1, "test2": db2},
shard_chooser=shard_chooser,
id_chooser=id_chooser,
execute_chooser=no_query_chooser,
)
return session
def test_selectinload_query(self):
session = ShardedSession(shards={"test": testing.db},
shard_chooser=lambda *args: 'test',
id_chooser=lambda *args: None,
query_chooser=lambda *args: ['test'])
Book, Page = self.classes("Book", "Page")
book = Book()
book.pages.append(Page())
session.add(book)
session.commit()
result = session.query(Book).options(selectinload('pages')).all()
eq_(result, [book])
def test_selectinload_query(self):
session = ShardedSession(
shards={"test": testing.db},
shard_chooser=lambda *args: 'test',
id_chooser=lambda *args: None,
query_chooser=lambda *args: ['test']
)
Book, Page = self.classes("Book", "Page")
book = Book()
book.pages.append(Page())
session.add(book)
session.commit()
result = session.query(Book).options(selectinload('pages')).all()
eq_(result, [book])
def test_lazy_load_from_identity_map(self):
session = ShardedSession(shards={"test": testing.db},
shard_chooser=lambda *args: 'test',
id_chooser=lambda *args: ['test'],
query_chooser=lambda *args: ['test'])
Book, Page = self.classes("Book", "Page")
book = Book()
book.pages.append(Page())
session.add(book)
session.commit()
book = session.query(Book).first()
page = session.query(Page).first()
def go():
eq_(page.book, book)
# doesn't emit SQL
self.assert_sql_count(testing.db, go, 0)
def execute_chooser(query: ORMExecuteState, ):
"""
For a given ORMExecuteState, returns the list of shard_ids where the query should be issued. Results from all shards returned will be combined together into a single listing.
"""
# stmt: Select = query.statement
# comparison = _get_query_comparisons(stmt)
chosen_shards = list(shards)
print("execute_chooser chosen", chosen_shards)
return chosen_shards
with ShardedSession(
shards=shards,
shard_chooser=shard_chooser,
id_chooser=id_chooser,
# execute_chooser=execute_chooser,
) as ses:
n_user = 100
user_list = [User(id=i) for i in range(1, n_user + 1)]
random.shuffle(user_list)
ses.add_all(user_list)
ses.commit()
# --- id_chooser
# get row by primary key will use id_chooser
print("=== id_chooser example ===")
user = ses.get(User, 1)
print(user)
# --- execute_chooser
