SQLAlchemy - можете ли вы добавить пользовательские методы к объекту запроса?

Есть ли способ создать пользовательские методы для объекта запроса, чтобы вы могли сделать что-то вроде этого?

User.query.all_active()

Где all_active() по существу .filter(User.is_active == True)

И уметь отфильтровывать его?

User.query.all_active().filter(User.age == 30)

Ответы

Ответ 1

Вы можете подклассифицировать базовый класс Query, чтобы добавить свои собственные методы:

from sqlalchemy.orm import Query

class MyQuery(Query):

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

Затем вы указываете SQLAlchemy использовать этот новый класс запросов при создании сеанса (docs здесь). Из вашего кода кажется, что вы можете использовать Flask-SQLAlchemy, поэтому вы сделали бы это следующим образом:

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

В противном случае вы передадите аргумент непосредственно в sessionmaker:

sessionmaker(bind=engine, query_cls=MyQuery)

Как сейчас, этот новый объект запроса не интересен, потому что мы строго кодировали класс User в методе, поэтому он не будет работать ни на что другое. Лучшая реализация будет использовать базовый класс запроса, чтобы определить, какой фильтр применить. Это немного сложно, но может быть сделано также:

class MyOtherQuery(Query):

  def _get_models(self):
    """Returns the query 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)

Наконец, этот новый класс запроса не будет использоваться динамическими отношениями (если они есть). Чтобы позволить им использовать его, вы можете передать его в качестве аргумента при создании отношения:

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

Ответ 2

Чтобы предоставить пользовательский метод, который будет использоваться всеми вашими моделями, которые наследуются от определенного родителя, сначала, как уже упоминалось, наследуйте от класса Query:

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)

Это для фляги-sqlalchemy, для которой люди все еще будут здесь, когда ищут этот ответ.