Function main

examples/sharding/separate_tables.py:265–334

()

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

        # 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",
        }

        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",
        }

        # 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]