mirror of
https://github.com/sqlalchemy/sqlalchemy.git
synced 2026-05-14 12:47:22 -04:00
Mike Bayer
8a32f36717
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
327 lines
10 KiB
Python
327 lines
10 KiB
Python
"""Illustrates sharding using distinct SQLite databases."""
|
|
from __future__ import annotations
|
|
|
|
import datetime
|
|
|
|
from sqlalchemy import Column
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy import ForeignKey
|
|
from sqlalchemy import inspect
|
|
from sqlalchemy import Integer
|
|
from sqlalchemy import select
|
|
from sqlalchemy import Table
|
|
from sqlalchemy.ext.horizontal_shard import set_shard_id
|
|
from sqlalchemy.ext.horizontal_shard import ShardedSession
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
from sqlalchemy.orm import Mapped
|
|
from sqlalchemy.orm import mapped_column
|
|
from sqlalchemy.orm import relationship
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.sql import operators
|
|
from sqlalchemy.sql import visitors
|
|
|
|
|
|
echo = True
|
|
db1 = create_engine("sqlite://", echo=echo)
|
|
db2 = create_engine("sqlite://", echo=echo)
|
|
db3 = create_engine("sqlite://", echo=echo)
|
|
db4 = create_engine("sqlite://", echo=echo)
|
|
|
|
|
|
# create session function. this binds the shard ids
|
|
# to databases within a ShardedSession and returns it.
|
|
Session = sessionmaker(
|
|
class_=ShardedSession,
|
|
shards={
|
|
"north_america": db1,
|
|
"asia": db2,
|
|
"europe": db3,
|
|
"south_america": db4,
|
|
},
|
|
)
|
|
|
|
|
|
# mappings and tables
|
|
class Base(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))
|
|
|
|
|
|
def id_generator(ctx):
|
|
# in reality, might want to use a separate transaction for this.
|
|
with db1.begin() as conn:
|
|
nextid = conn.scalar(ids.select().with_for_update())
|
|
conn.execute(ids.update().values({ids.c.nextid: ids.c.nextid + 1}))
|
|
return nextid
|
|
|
|
|
|
# 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.
|
|
|
|
|
|
class WeatherLocation(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 = continent
|
|
self.city = city
|
|
|
|
|
|
class Report(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
|
|
|
|
|
|
# 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",
|
|
}
|
|
|
|
|
|
def shard_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.
|
|
|
|
"""
|
|
if isinstance(instance, WeatherLocation):
|
|
return shard_lookup[instance.continent]
|
|
else:
|
|
return shard_chooser(mapper, instance.location)
|
|
|
|
|
|
def identity_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.
|
|
|
|
"""
|
|
if lazy_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"]
|
|
|
|
|
|
def execute_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 ids
|
|
for column, operator, value in _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'".
|
|
if column.shares_lineage(WeatherLocation.__table__.c.continent):
|
|
if operator == operators.eq:
|
|
ids.append(shard_lookup[value])
|
|
elif operator == operators.in_op:
|
|
ids.extend(shard_lookup[v] for v in value)
|
|
|
|
if len(ids) == 0:
|
|
return ["north_america", "asia", "europe", "south_america"]
|
|
else:
|
|
return ids
|
|
|
|
|
|
def _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 = []
|
|
|
|
def visit_bindparam(bind):
|
|
# visit a bind parameter.
|
|
|
|
value = bind.effective_value
|
|
binds[bind] = value
|
|
|
|
def visit_column(column):
|
|
clauses.add(column)
|
|
|
|
def visit_binary(binary):
|
|
if binary.left in clauses and binary.right in binds:
|
|
comparisons.append(
|
|
(binary.left, binary.operator, binds[binary.right])
|
|
)
|
|
|
|
elif binary.left in binds and binary.right in clauses:
|
|
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.
|
|
if statement.whereclause is not None:
|
|
visitors.traverse(
|
|
statement.whereclause,
|
|
{},
|
|
{
|
|
"bindparam": visit_bindparam,
|
|
"binary": visit_binary,
|
|
"column": visit_column,
|
|
},
|
|
)
|
|
return comparisons
|
|
|
|
|
|
# further configure create_session to use these functions
|
|
Session.configure(
|
|
shard_chooser=shard_chooser,
|
|
identity_chooser=identity_chooser,
|
|
execute_chooser=execute_chooser,
|
|
)
|
|
|
|
|
|
def setup():
|
|
# create tables
|
|
for db in (db1, db2, db3, db4):
|
|
Base.metadata.create_all(db)
|
|
|
|
# establish initial "id" in db1
|
|
with db1.begin() as conn:
|
|
conn.execute(ids.insert(), {"nextid": 1})
|
|
|
|
|
|
def main():
|
|
setup()
|
|
|
|
# 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))
|
|
|
|
with Session() as sess:
|
|
|
|
sess.add_all(
|
|
[tokyo, newyork, toronto, london, dublin, brasilia, quito]
|
|
)
|
|
|
|
sess.commit()
|
|
|
|
t = sess.get(WeatherLocation, tokyo.id)
|
|
assert t.city == tokyo.city
|
|
assert t.reports[0].temperature == 80.0
|
|
|
|
# select across shards
|
|
asia_and_europe = sess.execute(
|
|
select(WeatherLocation).filter(
|
|
WeatherLocation.continent.in_(["Europe", "Asia"])
|
|
)
|
|
).scalars()
|
|
|
|
assert {c.city for c in asia_and_europe} == {
|
|
"Tokyo",
|
|
"London",
|
|
"Dublin",
|
|
}
|
|
|
|
# optionally set a shard id for the query and all related loaders
|
|
north_american_cities_w_t = sess.execute(
|
|
select(WeatherLocation)
|
|
.filter(WeatherLocation.city.startswith("T"))
|
|
.options(set_shard_id("north_america"))
|
|
).scalars()
|
|
|
|
# Tokyo not included since not in the north_america shard
|
|
assert {c.city for c in north_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]
|
|
|
|
assert inspect(newyork_report).identity_key == (
|
|
Report,
|
|
(1,),
|
|
"north_america",
|
|
)
|
|
assert inspect(tokyo_report).identity_key == (Report, (1,), "asia")
|
|
|
|
# the token representing the originating shard is also available
|
|
# directly
|
|
assert inspect(newyork_report).identity_token == "north_america"
|
|
assert inspect(tokyo_report).identity_token == "asia"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|