0

Using Flask SQLAlchemy I'm querying a MySQL database with a table called bar and looking for rows that match a filter consisting of foo and country_code:

foo_filter = 'hello'
country_code_filter = 'ES'
result = Bar.filter_by(foo=foo_filter, country_code=country_code_filter).first()

The above code will return the first row in which foo = foo_filter and country_code = country_code_filter.

However, it's possible that we might not have rows for certain country codes that match foo. In these cases (i.e, cases where the above query returns 0 results), I'd like to use a default country filter of 'RoW' as our dataset should always have an RoW value for each possible value of foo. In the unexpected occurrence that this also doesn't return any results then an error should be thrown. This is the code I have for this:

foo_filter = 'hello'
country_code_filter = 'ES'
result = Bar.filter_by(foo=foo_filter, country_code=country_code_filter).first()
if not result:
    result = Bar.filter_by(foo=foo_filter, country_code='RoW').first()
if not result:
    raise RuntimeException(f"No data for combination {foo_filter}, {country_code_filter} or {foo_filter}, RoW")

This approach of running similar queries multiple times and checking the result each time until I get a row feels very messy/wrong but I haven't been able to find any better approaches that allow you to set an 'alternative' filter when your initial query returns 0 rows in Flask SQLAlchemy

Is there a cleaner approach to this?

SuperShoot
  • 9,880
  • 2
  • 38
  • 55
Duck Hunt Duo
  • 299
  • 1
  • 11

1 Answers1

1

This solution uses regular SQLAlchemy but running it on Flask-SQLAlchemy is just a couple of syntax changes away.

The idea is that we query for both the desired country code and the fall back one, ensure that the results are always sorted with the desired country code first and then limit the result to 1. E.g.:

result = (
    s.query(Bar)
    .filter(
        Bar.foo == foo_filter,
        Bar.country_code.in_([country_code_filter, "RoW"]),
    )
    .order_by(func.FIELD(Bar.country_code, country_code_filter, "RoW"))
    .limit(1)
    .first()
)

The FIELD function in MySQL allows you to specify a custom sort order, in this case it ensures that if a result exists with the country code you want, it will always be returned first. You can read more about FIELD here.

Here's the full test code:

from sqlalchemy_app import Base, Session, engine  # you need to create these
from sqlalchemy import Column, Integer, String, func


class Bar(Base):
    id = Column(Integer, primary_key=True)
    foo = Column(String(32))
    country_code = Column(String(32))


if __name__ == "__main__":
    Base.metadata.create_all(engine)
    s = Session()
    s.add_all(
        [
            Bar(foo="hello", country_code="ES"),
            Bar(foo="hello", country_code="ZIM"),
            Bar(foo="hello", country_code="RoW"),
        ]
    )
    s.commit()
    for foo_filter, country_code_filter, exp_res in (
        ("hello", "ES", "ES"),
        ("hello", "ZIM", "ZIM"),
        ("hello", "AUS", "RoW"),
        ("goodbye", "GBR", None),
    ):
        result = (
            s.query(Bar)
            .filter(
                Bar.foo == foo_filter,
                Bar.country_code.in_([country_code_filter, "RoW"]),
            )
            .order_by(func.FIELD(Bar.country_code, country_code_filter, "RoW"))
            .limit(1)
            .first()
        )
        assert getattr(result, "country_code", None) == exp_res

As you say that you want to raise an exception in the case that the query returns None, you can swap the .first() access method to .one() since the query limits the result set to 1 result server-side, you'll never hit the exception for returning more than 1 result and will raise a NoResultFound exception when no rows match the query.

SuperShoot
  • 9,880
  • 2
  • 38
  • 55