"""Illustrates sharding API used with asyncio.For the sync version of this example, see separate_databases.py.Most of the code here is copied from separate_databases.py and worksin exactly the same way. The main change is how the``async_sessionmaker`` is configured, and as is specific to this examplethe routine that generates new primary keys."""from__future__importannotationsimportasyncioimportdatetimefromsqlalchemyimportColumnfromsqlalchemyimportForeignKeyfromsqlalchemyimportinspectfromsqlalchemyimportIntegerfromsqlalchemyimportselectfromsqlalchemyimportTablefromsqlalchemy.ext.asyncioimportasync_sessionmakerfromsqlalchemy.ext.asyncioimportcreate_async_enginefromsqlalchemy.ext.horizontal_shardimportset_shard_idfromsqlalchemy.ext.horizontal_shardimportShardedSessionfromsqlalchemy.ormimportDeclarativeBasefromsqlalchemy.ormimportimmediateloadfromsqlalchemy.ormimportMappedfromsqlalchemy.ormimportmapped_columnfromsqlalchemy.ormimportrelationshipfromsqlalchemy.sqlimportoperatorsfromsqlalchemy.sqlimportvisitorsecho=Truedb1=create_async_engine("sqlite+aiosqlite://",echo=echo)db2=create_async_engine("sqlite+aiosqlite://",echo=echo)db3=create_async_engine("sqlite+aiosqlite://",echo=echo)db4=create_async_engine("sqlite+aiosqlite://",echo=echo)# for asyncio, the ShardedSession class is passed# via sync_session_class. The shards themselves are used within# implicit-awaited internals, so we use the sync_engine Engine objects# in the shards dictionary.Session=async_sessionmaker(sync_session_class=ShardedSession,expire_on_commit=False,shards={"north_america":db1.sync_engine,"asia":db2.sync_engine,"europe":db3.sync_engine,"south_america":db4.sync_engine,},)# mappings and tablesclassBase(DeclarativeBase):pass# we need a way to create identifiers which are unique across all databases.# one easy way would be to just use a composite primary key, where one value# is the shard id. but here, we'll show something more "generic", an id# generation function. we'll use a simplistic "id table" stored in database# #1. Any other method will do just as well; UUID, hilo, application-specific,# etc.ids=Table("ids",Base.metadata,Column("nextid",Integer,nullable=False))defid_generator(ctx):# id_generator is run within a "synchronous" context, where# we use an implicit-await API that will convert back to explicit await# calls when it reaches the driver.withdb1.sync_engine.begin()asconn:nextid=conn.scalar(ids.select().with_for_update())conn.execute(ids.update().values({ids.c.nextid:ids.c.nextid+1}))returnnextid# table setup. we'll store a lead table of continents/cities, and a secondary# table storing locations. a particular row will be placed in the database# whose shard id corresponds to the 'continent'. in this setup, secondary rows# in 'weather_reports' will be placed in the same DB as that of the parent, but# this can be changed if you're willing to write more complex sharding# functions.classWeatherLocation(Base):__tablename__="weather_locations"id:Mapped[int]=mapped_column(primary_key=True,default=id_generator)continent:Mapped[str]city:Mapped[str]reports:Mapped[list[Report]]=relationship(back_populates="location")def__init__(self,continent:str,city:str):self.continent=continentself.city=cityclassReport(Base):__tablename__="weather_reports"id:Mapped[int]=mapped_column(primary_key=True)location_id:Mapped[int]=mapped_column(ForeignKey("weather_locations.id"))temperature:Mapped[float]report_time:Mapped[datetime.datetime]=mapped_column(default=datetime.datetime.now)location:Mapped[WeatherLocation]=relationship(back_populates="reports")def__init__(self,temperature:float):self.temperature=temperature# step 5. define sharding functions.# we'll use a straight mapping of a particular set of "country"# attributes to shard id.shard_lookup={"North America":"north_america","Asia":"asia","Europe":"europe","South America":"south_america",}defshard_chooser(mapper,instance,clause=None):"""shard chooser. looks at the given instance and returns a shard id note that we need to define conditions for the WeatherLocation class, as well as our secondary Report class which will point back to its WeatherLocation via its 'location' attribute. """ifisinstance(instance,WeatherLocation):returnshard_lookup[instance.continent]else:returnshard_chooser(mapper,instance.location)defidentity_chooser(mapper,primary_key,*,lazy_loaded_from,**kw):"""identity chooser. given a primary key, returns a list of shards to search. here, we don't have any particular information from a pk so we just return all shard ids. often, you'd want to do some kind of round-robin strategy here so that requests are evenly distributed among DBs. """iflazy_loaded_from:# if we are in a lazy load, we can look at the parent object# and limit our search to that same shard, assuming that's how we've# set things up.return[lazy_loaded_from.identity_token]else:return["north_america","asia","europe","south_america"]defexecute_chooser(context):"""statement execution chooser. this also returns a list of shard ids, which can just be all of them. but here we'll search into the execution context in order to try to narrow down the list of shards to SELECT. """ids=[]# we'll grab continent names as we find them# and convert to shard idsforcolumn,operator,valuein_get_select_comparisons(context.statement):# "shares_lineage()" returns True if both columns refer to the same# statement column, adjusting for any annotations present.# (an annotation is an internal clone of a Column object# and occur when using ORM-mapped attributes like# "WeatherLocation.continent"). A simpler comparison, though less# accurate, would be "column.key == 'continent'".ifcolumn.shares_lineage(WeatherLocation.__table__.c.continent):ifoperator==operators.eq:ids.append(shard_lookup[value])elifoperator==operators.in_op:ids.extend(shard_lookup[v]forvinvalue)iflen(ids)==0:return["north_america","asia","europe","south_america"]else:returnidsdef_get_select_comparisons(statement):"""Search a Select or Query object for binary expressions. Returns expressions which match a Column against one or more literal values as a list of tuples of the form (column, operator, values). "values" is a single value or tuple of values depending on the operator. """binds={}clauses=set()comparisons=[]defvisit_bindparam(bind):# visit a bind parameter.value=bind.effective_valuebinds[bind]=valuedefvisit_column(column):clauses.add(column)defvisit_binary(binary):ifbinary.leftinclausesandbinary.rightinbinds:comparisons.append((binary.left,binary.operator,binds[binary.right]))elifbinary.leftinbindsandbinary.rightinclauses:comparisons.append((binary.right,binary.operator,binds[binary.left]))# here we will traverse through the query's criterion, searching# for SQL constructs. We will place simple column comparisons# into a list.ifstatement.whereclauseisnotNone:visitors.traverse(statement.whereclause,{},{"bindparam":visit_bindparam,"binary":visit_binary,"column":visit_column,},)returncomparisons# further configure create_session to use these functionsSession.configure(shard_chooser=shard_chooser,identity_chooser=identity_chooser,execute_chooser=execute_chooser,)asyncdefsetup():# create tablesfordbin(db1,db2,db3,db4):asyncwithdb.begin()asconn:awaitconn.run_sync(Base.metadata.create_all)# establish initial "id" in db1asyncwithdb1.begin()asconn:awaitconn.execute(ids.insert(),{"nextid":1})asyncdefmain():awaitsetup()# save and load objects!tokyo=WeatherLocation("Asia","Tokyo")newyork=WeatherLocation("North America","New York")toronto=WeatherLocation("North America","Toronto")london=WeatherLocation("Europe","London")dublin=WeatherLocation("Europe","Dublin")brasilia=WeatherLocation("South America","Brasila")quito=WeatherLocation("South America","Quito")tokyo.reports.append(Report(80.0))newyork.reports.append(Report(75))quito.reports.append(Report(85))asyncwithSession()assess:sess.add_all([tokyo,newyork,toronto,london,dublin,brasilia,quito])awaitsess.commit()t=awaitsess.get(WeatherLocation,tokyo.id,options=[immediateload(WeatherLocation.reports)],)assertt.city==tokyo.cityassertt.reports[0].temperature==80.0# select across shardsasia_and_europe=(awaitsess.execute(select(WeatherLocation).filter(WeatherLocation.continent.in_(["Europe","Asia"])))).scalars()assert{c.cityforcinasia_and_europe}=={"Tokyo","London","Dublin",}# optionally set a shard id for the query and all related loadersnorth_american_cities_w_t=(awaitsess.execute(select(WeatherLocation).filter(WeatherLocation.city.startswith("T")).options(set_shard_id("north_america")))).scalars()# Tokyo not included since not in the north_america shardassert{c.cityforcinnorth_american_cities_w_t}=={"Toronto",}# the Report class uses a simple integer primary key. So across two# databases, a primary key will be repeated. The "identity_token"# tracks in memory that these two identical primary keys are local to# different shards.newyork_report=newyork.reports[0]tokyo_report=tokyo.reports[0]assertinspect(newyork_report).identity_key==(Report,(1,),"north_america",)assertinspect(tokyo_report).identity_key==(Report,(1,),"asia")# the token representing the originating shard is also available# directlyassertinspect(newyork_report).identity_token=="north_america"assertinspect(tokyo_report).identity_token=="asia"if__name__=="__main__":asyncio.run(main())