Querying multiple scalar values in SQLAlchemy

SQLAlchemy is probably the greatest ORM library ever invented — and not just for Python. Like everything, though, it could sometimes use an improvement.

One typical pattern that I’ve seen repeated many times is querying for values of a single column from multiple rows. Almost always, the desired outcome is to have those values in a Python iterable. Unfortunately, when returning multiple rows, the Query class can only give them out as tuples, making it necessary to unpack them explicitly every time:

active_users_ids = [
    uid for (uid,) in session.query(User.id).filter(User.is_active).all()
]

It’s worth noting that a very similar operation — retrieving a single column value from one row — is comparatively cleaner thanks to a dedicated Query.scalar method. So why not have a similar one for the use case above?…

Turns out that SQLAlchemy is so amazingly extensible that we can easily outfit Query with such a method ourselves.

Our own Query

The officially ordained way to do something like that is to extend the original sqlalchemy.orm.query.Query class with any additional functionality we require. We can then introduce a method in the subclass — let’s call it values — that implements our logic:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

from sqlalchemy.orm.exc import MultipleResultsFound
from sqlalchemy.orm.query import Query as _Query


class Query(_Query):
    """Custom SQLAlchemy query class."""

    def values(self):
        """Return an iterable of all scalar, singleton element values
        from rows matched by this query.

        :raise MultipleValuesFound: If result rows have more than one element
        """
        try:
            return [x for (x,) in self.all()]
        except ValueError as e:
            raise MultipleValuesFound(str(e))


class MultipleValuesFound(ValueError, MultipleResultsFound):
    """Exception raised by :meth:`Query.values`
    when multiple values were found in a single result row.
    """

There is a special erroneous condition that occurrs when the Query returns rows with more than one column value. It would be possible to detect it preemptively by inspecting certain attributes of the query object, but it’s simpler to just catch the exception caused by unsuccessful tuple unpacking.

In any case, I recommend re-raising it as a custom exception to enable more robust error handling in the client code.

Wiring it

The last piece of the puzzle is to make Session objects in our application use the new Query class. It’s easy: we simply pass it as query_cls parameter to the Session constructor. In more realistic scenarios — such as web applications with per-request database sessions — it means giving it directly to the sessionmaker:

from sqlalchemy import create_engine
from sqlalchemy import scoped_session, sessionmaker

from myapp.ext.sqlalchemy import Query


session = scoped_session(sessionmaker(bind=create_engine(DATABASE_URL),
                                      query_cls=Query))

Of course, actual session setup in your app may differ in few details from the sample above.

Usage example

With the new Query class wired in, the values method becomes immediately usable:

active_user_ids = session.query(User.id).filter(User.is_active).values()

and the ugly manual unpacking of row tuples is no longer necessary.