In SQLAlchemy, how do I query composite primary keys?

Question

I'm using SQLAlchemy to programmatically query a table with a composite foreign key. e.g.:

CREATE TABLE example (
id INT NOT NULL,
date TIMESTAMP NOT NULL,
data VARCHAR(128)
PRIMARY KEY (id, date)
)

I'd like to take a list of values and get rows back e.g.:

interesting_data = (
    (1, '2016-5-1'), 
    (1, '2016-6-1'),
    (2, '2016-6-1'),
    (3, '2016-5-1'),
    (3, '2016-6-1'),
)
select(
    [example.c.id, example.c.date, example.c.data],
).where(example.primary_key.in_(interesting_data)

If each column was independent I could do

interesting_ids = [1,2,3]
interesting_dates = ['2016-5-1', '2016-6-1']
select(
    [example.c.id, example.c.date, example.c.data],
).where(
    example.c.id.in_(interesting_ids)
).where(
    example.c.date.in_(interesting_dates)
)

But that clearly fails to bring only the unique matching (id, date) tuples. I suspect there is a way to specify the compound primary key to query, but I can't find any documentation after a search.

Arsenic

Answer 1

Assuming your model class is called Example with (id, date) is the composite primary key:

You can query either with:

import sqlalchemy
...
Example.query.get((id, date))

or

import sqlalchemy
from sqlalchemy.orm import sessionmaker
...
engine = sqlalchemy.create_engine('postgresql://user:pass@localhost/db_name')
session = sessionmaker(bind=engine)()
session.query(Example).get((id, date))

Incorrupt

You can also simply do session.get(<TABLE>, (value1, value2)). – Placidia

Answer 2

Use a list comprehension in your where clause:

from sqlalchemy import and_, or_, select

stmt = select(
    [example.c.id, example.c.date, example.c.data]
).where(or_(and_(example.c.id==data[0], example.c.date==data[1])
            for data in interesting_data))

However, a separate issue I noticed is that you're comparing a date column to string data type. The interesting_data list should be

import datetime as dt

interesting_data = (
    (1, dt.date(2016,5,1)), 
    (1, dt.date(2016,6,1)),
    (2, dt.date(2016,6,1)),
    (3, dt.date(2016,6,1)),
    (3, dt.date(2016,6,1)),
)

Also, note that it is possible to create a basic statement, and then add clauses to it incrementally, leading to (hopefully) better legibility and code reuse.

So, it would be possible to write the above as

base_stmt = select([example.c.id, example.c.date, example.c.data])
wheres = or_(and_(example.c.id==data[0], example.c.date==data[1])
             for data in interesting_data))
stmt = base_stmt.where(wheres)

This generates the following sql (new-lines & spaces added by me):

SELECT example.id, example.date, example.data 
FROM example 
WHERE 
   example.id = :id_1 AND example.date = :date_1 
OR example.id = :id_2 AND example.date = :date_2
OR example.id = :id_3 AND example.date = :date_3 
OR example.id = :id_4 AND example.date = :date_4 
OR example.id = :id_5 AND example.date = :date_5

note: if you have a lot of rows to filter like this, it may be more efficient to create a temporary table, insert rows into this temporary table from interesting_data, and then inner join to this table, rather than add the where clause as shown above.

Heir