Complex joins with Peewee

Question

It's rather embarrassing to ask this question, since it seems to be so trivial, yet I can't find a solution that works.

I have the following function:

def planner(departure_id, arrival_id):
    departure = Stop.get(Stop.id == departure_id)
    arrival = Stop.get(Stop.id == arrival_id)
    buses = Bus.select().join(RideStopRelationship).join(Stop).where(Stop.id == departure)
    for bus in buses:
        print bus.line
        for stop in bus.stops:
            print stop.time, stop.stop.name

based on the following models:

class Stop(BaseModel):
    name = CharField()
    #lat = FloatField()
    #lng = FloatField()

class Bus(BaseModel):
    line = IntegerField()
    number = IntegerField()
    direction = IntegerField()

class RideStopRelationship(BaseModel):
    bus = ForeignKeyField(Bus, related_name = "stops")
    stop = ForeignKeyField(Stop, related_name = "buses")
    time = TimeField()

The crucial line is Bus.select().join(RideStopRelationship).join(Stop).where(Stop.id == departure). I'm trying to get all buses that will stop at both departure and arrival. However, the above query returns all buses that stop at departure. How would I get buses that stop at both 'departure' and 'arrival'?

If I'm making this too complicated (either my models being too complicated, or my query), feel free to correct me.

EDIT: There's one way that does work:

buses_departure = Bus.select().join(RideStopRelationship).join(Stop).where(Stop.id == departure)
buses_arrival = Bus.select().join(RideStopRelationship).join(Stop).where(Stop.id == arrival)
buses = Bus.select().where(Bus.id << buses_departure & Bus.id << buses_arrival)

but it's rather long for what should be a simply query...

Answer 1

You might try something like this:

departure = Stop.get(...)
arrival = Stop.get(...)
query = (Bus
         .select(Bus)
         .join(RideStopRelationship)
         .where(RideStopRelationship.stop << [departure, arrival])
         .group_by(Bus)
         .having(fn.Count(Bus.id) == 2))

Unrelated, but one thing to note is that due to the way python evaluates operators, you need to put parentheses around your in queries:

buses = Bus.select().where(
    (Bus.id << buses_departure) & 
    *Bus.id << buses_arrival))