SQLAlchemy: è possibile aggiungere metodi personalizzati all'oggetto query?

Question

C'è un modo per creare metodi personalizzati per l'oggetto query in modo che tu possa fare qualcosa di simile?

User.query.all_active() 

Dove all_active() è essenzialmente .filter(User.is_active == True)

Ed essere in grado di filtrare fuori di esso?

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

Answer 1

È possibile creare una sottoclasse la base Query classe per aggiungere i propri metodi:

from sqlalchemy.orm import Query 

class MyQuery(Query): 

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

È quindi dire SQLAlchemy per utilizzare questa nuova classe di query quando si crea la sessione (docs here). Dal codice sembra che si potrebbero utilizzare Flask-SQLAlchemy, così si avrebbe fatto come segue:

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

altrimenti si sarebbe passare l'argomento direttamente al sessionmaker:

sessionmaker(bind=engine, query_cls=MyQuery) 

Come del diritto ora, questo nuovo oggetto query non è così interessante perché abbiamo codificato con hardcoded la classe User nel metodo, quindi non funzionerà per nient'altro. Un'implementazione migliore utilizza la classe sottostante della query per determinare quale filtro applicare. Questo è un po 'complicato, ma può essere fatto così:

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) 

Infine, questa nuova classe di query non saranno utilizzati da relazioni dinamiche (se ne avete). Per permettere a coloro usare anche esso, è possibile passarlo come argomento quando si crea la relazione:

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