35

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_active() is essentially .filter(User.is_active == True)

And be able to filter off of it?

User.query.all_active().filter(User.age == 30)
mphuie
  • 970
  • 1
  • 9
  • 17
  • I wonder if anything in this other [question/answer](http://stackoverflow.com/questions/7604967/sqlalchemy-build-query-filter-dynamically-from-dict) could be of any use for this? (It's less of a custom method... and more of a custom dynamic filter, though.) – summea Apr 10 '13 at 20:55

3 Answers3

60

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)
mtth
  • 4,671
  • 3
  • 30
  • 36
  • This doesn't work for me. https://gist.github.com/nickretallack/76bcd65cc0305abcd33b – Nick Retallack Oct 22 '14 at 23:07
  • 1
    @NickRetallack I had some success with creating a BaseModel class and declaring the query_class class attribute. However, this seems to be the issue: https://github.com/mitsuhiko/flask-sqlalchemy/blob/fe67c633f2e6d66a01a1670e5fc6649506358d20/flask_sqlalchemy/__init__.py#L737 – justanr May 05 '15 at 13:53
  • 4
    So you say that one session can work only with one query class. Can we do better? and allow the same session to choose dynamically the query class? I mean different tables can have different query classes, but I may want to use them in the same session (for example when using them in one transaction). – Mahdi Jul 28 '16 at 08:36
  • 1
    I think there is a typo, shouln't it be `query._column_descriptions` (with the underscore)? Also maybe add the imports, would help, even if it seems to be all from the orm package. – Brizar Jul 11 '21 at 22:31
2

this work for me finely

from sqlalchemy.orm import query
from flask_sqlalchemy import BaseQuery

    class ParentQuery(BaseQuery):
        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')))
    
    #using like this
    class AccountWorkflowModel(db.Model):
        query_class = ParentQuery
        .................

Second way in ver 2.0

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Query

# Create the SQLAlchemy engine and session
engine = create_engine('your_database_url')
Session = sessionmaker(bind=engine)
session = Session()


# Define your custom query subclass
class CustomQuery(Query):
    def custom_method(self, param):
        # Implement your custom query method logic here
        return self.filter_by(column=param)


# Associate the custom query subclass with the session
session.query = session.query_property(CustomQuery)


# Now you can use your custom method on the query object
results = session.query(Model).custom_method('your_param').all()
Shahram
  • 27
  • 4
  • Can you please show complete code of this. I want to add custom filter in sqlalchemy but I also want it get initialized in Related model ( to code in model class only and not want to touch the sqlalchemy conf like session and engine etc). I want to write this kind of code https://stackoverflow.com/a/56806069/5470263 it seems like it can be achived with your approach. – danish May 21 '20 at 07:53
  • hi, this is complete code. where you have problem? please tell me, i could send the specify code that you need. usage : @staticmethod def getAllItems(): return AccountWorkflowModel.query.FilterByCustomer().all() – Shahram May 22 '20 at 08:32
  • @Shahram where does the variable `query` on line no 3 come from??? or did u mean to write `self` – lordvcs Nov 08 '22 at 05:54
  • I imported from sqlalchemy.orm import query – Shahram Nov 09 '22 at 06:33
  • Is there another way for SQLAlchemy 2.0? It is considered legacy now, even though Query will remain. I would like this to be more like v2.0 – Miradil Zeynalli Aug 27 '23 at 13:39
  • I think there are a little way for SQLAlchemy, I update my answer with new approach for v 2.0 – Shahram Aug 29 '23 at 04:31
1

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.