SQLAlchemy: Custom Sorting Function with Arguments

Recently solved a problem, that wasn’t in bing results. AI suggestions about SQLAlchemy’s hybrid properties and expressions were misleading, because they can’t take external arguments. But what if you want to sort your model on a computed value, that depends on column(s) and external argument?

Real World Problem

Product model:

class Product(Base):
    __tablename__ = 'product'

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(40))
    category: Mapped[str] = mapped_column(String(40))
    price: Mapped[Decimal] = mapped_column(Numeric(precision=10, scale=2))

    def client_price(self, cl_disc_config: dict[str, str]) -> Decimal:
        '''returns client price with computed client discount'''
        cat_disc_str: dict = cl_disc_config.get(self.category, '0')
        disc_fraction = Decimal(cat_disc_str) / 100
        return round(self.price * (1 - disc_fraction), 2)

    def __repr__(self):
        return f'Product(name={self.name}, category={self.category}, price={self.price})'

Clients make Orders, orders have OrderItem’s, the usual stuff. The catch here is that some clients have Product category-based discounts defined. And when they view product catalog, prices should be tailed based on their discounts.

Minimal client model:

class Client(Base):
    __tablename__ = 'client'

    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(60))
    disc_config: Mapped[Optional[dict | list]] = mapped_column(type_=JSON)

    def __repr__(self):
        return f'Client(name={self.name}, email={self.email})'

Perhaps there are better ways to represent this in some other database schema, but this was the route I chose. It worked well, until I had a paginated (filterable, sortable) catalog and had to implement sorting on not Product.price, but Product.client_price(client.disc_config), that takes named discount configuration.

Here is a disc_config column value example:

{'electronics': '10', 'furniture': '5', 'drug': '25', 'other': '0'}

Custom Sorting Function

For front end, at least with jinja2, you can write whole expression, that would render client price in html:

{% for p in catalog %}
  {{ p.client_price(client.disc_config) }}
{% endfor %}

But you can’t apply that to query statement, without being an animal iterating over…

Here’s a custom sort function, that takes statement, disc_config, sort direction and returns statement back:

from sqlalchemy import Select, case, asc, desc

def apply_client_price_sort(stmt: Select, disc_config: dict[str, str] | None, sort_asc: bool = True) -> Select:
    '''applies custom sort on computed client price based on client disc_config
    Args:
        - stmt - sqlalchemy select statement for SKU
        - disc_config - client discount configuration
        - sort_asc - sort direction'''
    if disc_config is None:
        return stmt.order_by(asc(Product.price)) if sort_asc else stmt.order_by(desc(Product.price))
    cases = []
    for category, disc_str in disc_config.items():
        if disc_str != '0':
            disc_fraction = Decimal(disc_str) / 100
            cases.append((Product.category == category, Product.price * (1 - disc_fraction)))
    clprice_expr = case(*cases, else_=Product.price)
    sort_order = asc(clprice_expr) if sort_asc else desc(clprice_expr)
    return stmt.order_by(sort_order)

In case client discount configuration is undefined (client.disc_config is None), or some new Product category has no corresponding configuration (missing dict key), we fall back to Product.price.

Applying Custom Sort Function

Pass query, to return sorted one. ProdCat below is just enum class.

    sober_catalog_stmt = select(Product).where(Product.category != ProdCat.DRUG.value)
    sober_catalog_stmt = apply_client_price_sort(sober_catalog_stmt, disc_config, sort_asc=False)
    catalog: list[Product] = session.scalars(sober_catalog_stmt).all()

    for p in catalog:
        cl_disc = disc_config.get(p.category, '0')
        cl_price = p.client_price(disc_config)
        print(f'{p.name} ({p.category}) price: {p.price}, -{cl_disc}% = {cl_price}')

Which, with sample products prints:

prod_13 (other) price: 85.24 -0% = 85.24
prod_5 (electronics) price: 88.00 -10% = 79.20
prod_16 (furniture) price: 79.37 -5% = 75.40
prod_12 (other) price: 73.99 -0% = 73.99
prod_0 (furniture) price: 72.54 -5% = 68.91
prod_10 (other) price: 46.95 -0% = 46.95
prod_3 (electronics) price: 37.51 -10% = 33.76
prod_2 (electronics) price: 28.04 -10% = 25.24
prod_15 (other) price: 20.24 -0% = 20.24
prod_7 (electronics) price: 16.89 -10% = 15.20
prod_6 (other) price: 14.44 -0% = 14.44
prod_1 (other) price: 11.23 -0% = 11.23

Complete Example

You can find complete standalone script example with Sqlite and SQLAlchemy in this gist. I define models, create tables, add sample random data and pull sorted catalog for specific client based on his discount configuration from most expensive to cheapest.