SQLAlchemy - can you add custom methods to the query object?

Question

Is there a way to create custom methods to the query object so you can do something like this?

User.query.all_active()

Where all_acti

Answer 1

this work for me finely

class ParentQuery(Query):
    def _get_models(self):     
        if hasattr(query, 'attr'):
            return [query.attr.target_mapper]
        else:
            return self._mapper_zero().class_

    def FilterByCustomer(self):
        model_class = self._get_models()
        return self.filter(model_class.customerId == int(g.customer.get('customerId')))


class AccountWorkflowModel(db.Model):
    query_class = ParentQuery
    .................

Answer 2

You can subclass the base Query class to add your own methods:

from sqlalchemy.orm import Query

class MyQuery(Query):

  def all_active(self):
    return self.filter(User.is_active == True)

You then tell SQLAlchemy to use this new query class when you create the session (docs here). From your code it looks like you might be using Flask-SQLAlchemy, so you would do it as follows:

db = SQLAlchemy(session_options={'query_cls': MyQuery})

Otherwise you would pass the argument directly to the sessionmaker:

sessionmaker(bind=engine, query_cls=MyQuery)

As of right now, this new query object isn't that interesting because we hardcoded the User class in the method, so it won't work for anything else. A better implementation would use the query's underlying class to determine which filter to apply. This is slightly tricky but can be done as well:

class MyOtherQuery(Query):

  def _get_models(self):
    """Returns the query's underlying model classes."""
    if hasattr(query, 'attr'):
      # we are dealing with a subquery
      return [query.attr.target_mapper]
    else:
      return [
        d['expr'].class_
        for d in query.column_descriptions
        if isinstance(d['expr'], Mapper)
      ]

  def all_active(self):
    model_class = self._get_models()[0]
    return self.filter(model_class.is_active == True)

Finally, this new query class won't be used by dynamic relationships (if you have any). To let those also use it, you can pass it as argument when you create the relationship:

users = relationship(..., query_class=MyOtherQuery)

Answer 3

To provide a custom method that will be used by all your models that inherit from a particular parent, first as mentioned before inherit from the Query class:

from flask_sqlalchemy import SQLAlchemy, BaseQuery
from sqlalchemy.inspection import inspect

class MyCustomQuery(BaseQuery):
    def all_active(self):
        # get the class
        modelClass = self._mapper_zero().class_
        # get the primary key column
        ins = inspect(modelClass)
        # get a list of passing objects
        passingObjs = []
        for modelObj in self:
            if modelObj.is_active == True:
                # add to passing object list
                passingObjs.append(modelObj.__dict__[ins.primary_key[0].name])
        # change to tuple
        passingObjs = tuple(passingObjs)
        # run a filter on the query object
        return self.filter(ins.primary_key[0].in_(passingObjs))

# add this to the constructor for your DB object
myDB = SQLAlchemy(query_class=MyCustomQuery)

This is for flask-sqlalchemy, for which people will still get here when looking for this answer.