CODEWITHSUNDEEP
sqlalchemyflask-sqlalchemy
Jason Strimpel
Jason Strimpel

Reputation: 15526

SQLAlchemy query construction

I'm working with a Key model backed by Postgres that is a generic table to hold API keys:

class Key(Model):
    __tablename__ = "keys"
    id = Column(Integer, primarykey=True)
    user_id = Column(Integer, ForeignKey("users.id"))
    brokerage_id = Column(Integer, ForeignKey("brokerages.id"))
    account_id = Column(Integer, ForeignKey("accounts.id"))
    key = Column(String(128))
    value = Column(String(128))

In the below example, user 2 has three keys. All three are associated with brokerage 2 and account 2. This is represented by IDs 4 through 6. For this site, the user has an authentication token plus two query IDs.

id  user_id brokerage_id    account_id      key         value
--------------------------------------------------------------------
4   2       2               2               token       999999999999
5   2       2               2               query_id    888888      
6   2       2               2               query_id    777777      
7   1       2               3               token       444444444444

I am trying to construct a query so that my result will be modeled as such:

[(user_id, brokerage_id, account_id, token, [query_id_1, query_id_2, ...]), ...]

So for the above example, it would look like this

[(2, 2, 2, 999999999999, [888888, 777777]), (1, 2, 3, 444444444444, [])]

I've got the following queries which select the token and the query_ids

tokens = db.session.query(
    Key.user_id, Key.brokerage_id, Key.account_id, Key.value
).filter(Key.key=='token').all()

query_ids = db.session.query(
    Key.user_id, Key.brokerage_id, Key.account_id, Key.value
).filter(Key.key=='query_id').all()

I've tried using subquery in various ways but cannot quite get the output to resemble what I need. How can I construct a query to return results in a way that align to my list of tuples, above?

Result

Adding the final working query here thanks to @rfkortekaas

from sqlalchemy.orm import aliased
from sqlalchemy import func, and_

from project.models import Key
from project.extensions import db

key_token = aliased(Key)

q = db.session.query(
    key_token.user_id,
    key_token.brokerage_id,
    key_token.account_id,
    key_token.value.label('token'),
    func.array_agg(Key.value).label('query_ids')
).join(
    Key,
    and_(
        key_token.user_id == Key.user_id,
        key_token.brokerage_id == Key.brokerage_id,
        key_token.account_id == Key.account_id,
        Key.key == 'query_id'
    )
).filter(
    key_token.key == 'token'
).group_by(
    key_token.user_id,
    key_token.brokerage_id,
    key_token.account_id,
    key_token.value
)
results = q.all()

Upvotes: 0

Views: 354

Answers (1)

rfkortekaas
rfkortekaas

Reputation: 6514

You can use the array_agg function from PostgreSQL to create an array of the results:

from sqlalchemy.orm import aliased

key_token = aliased(Key)

stmt = select(key_token.user_id,
              key_token.brokerage_id,
              key_token.account_id,
              key_token.value.label('token'),
              func.array_agg(Key.value).label('query_ids')
              ).join(Key,
                     and_(key_token.user_id == Key.user_id,
                          key_token.brokerage_id == Key.brokerage_id,
                          key_token.account_id == Key.account_id,
                          Key.key == 'query_id'))\
               .where(key_token.key == 'token')\
               .group_by(key_token.user_id,
                         key_token.brokerage_id,
                         key_token.account_id,
                         key_token.value)
keys = session.execute(stmt).all()
for row in keys:
    print(row)

Results in:

user_id brokerage_id account_id token query_ids
1 2 3 '44' ['4']
2 2 1 '33" ['6']
2 2 2 '99" ['8', '7]

For the following dataset:

user_id brokerage_id account_id key value
2 2 2 token '99'
2 2 1 token '33'
2 2 1 query_id '6'
2 2 2 query_id '8'
2 2 2 query_id '7'
1 2 3 token '44'
1 2 3 query_id '4'

Upvotes: 2

Related Questions