add set_shard_id() loader option for horizontal shard

Added new option to horizontal sharding API
:class:`_horizontal.set_shard_id` which sets the effective shard identifier
to query against, for both the primary query as well as for all secondary
loaders including relationship eager loaders as well as relationship and
column lazy loaders.

Modernize sharding examples with new-style mappings, add new asyncio example.

Fixes: #7226
Fixes: #7028
Change-Id: Ie69248060c305e8de04f75a529949777944ad511

1 files changed, 12 insertions, 11 deletions

diff --git a/lib/sqlalchemy/testing/assertions.py b/lib/sqlalchemy/testing/assertions.py
index 5dd32e8cc..c66ba71c3 100644
--- a/lib/sqlalchemy/testing/assertions.py
+++ b/lib/sqlalchemy/testing/assertions.py
@@ -906,7 +906,17 @@ class AssertsExecutionResults:
             db, callable_, assertsql.CountStatements(count)
         )
 
-    def assert_multiple_sql_count(self, dbs, callable_, counts):
+    @contextlib.contextmanager
+    def assert_execution(self, db, *rules):
+        with self.sql_execution_asserter(db) as asserter:
+            yield
+        asserter.assert_(*rules)
+
+    def assert_statement_count(self, db, count):
+        return self.assert_execution(db, assertsql.CountStatements(count))
+
+    @contextlib.contextmanager
+    def assert_statement_count_multi_db(self, dbs, counts):
         recs = [
             (self.sql_execution_asserter(db), db, count)
             for (db, count) in zip(dbs, counts)
@@ -915,21 +925,12 @@ class AssertsExecutionResults:
         for ctx, db, count in recs:
             asserters.append(ctx.__enter__())
         try:
-            return callable_()
+            yield
         finally:
             for asserter, (ctx, db, count) in zip(asserters, recs):
                 ctx.__exit__(None, None, None)
                 asserter.assert_(assertsql.CountStatements(count))
 
-    @contextlib.contextmanager
-    def assert_execution(self, db, *rules):
-        with self.sql_execution_asserter(db) as asserter:
-            yield
-        asserter.assert_(*rules)
-
-    def assert_statement_count(self, db, count):
-        return self.assert_execution(db, assertsql.CountStatements(count))
-
 
 class ComparesIndexes:
     def compare_table_index_with_expected(