Search code examples
pythonmongoengine

How to override mongoengine's QuerySet method?


How can I override a mongoengine's queryset method?

Specifically, I want to override .order_by(), but the closest I can get is to add another method ordered that would conditionally call .order_by():

class TransactionQuerySet(QuerySet):

    def ordered(self, *args, target_data=None, **kwargs):
        if target_data is None:
            result = self.order_by(*args, **kwargs)
            return result
        else:
            return 'transactions sorted by target data'

Ideally I would like this new method to be named the same as mongoengine's method - order_by - but if I do this I will exceed recursion depth when the queryset manager is called like Transaction.objects.order_by.

How can I do this?


Solution

  • To avoid recursion, you should call order_by method of the parent class. TransactionQuerySet should look like the following.

    class TransactionQuerySet(QuerySet):
    
        def order_by(self, *args, target_data=None):
            if target_data is None:
                return super().order_by(*args)
            return "transactions sorted by target data"
    

    Now if you call order_by on TransactionQuerySet object it won't fall in recursion.

    class Transaction(Document):
        title = StringField(max_length=100, required=True)
        last_modified = DateTimeField(default=datetime.utcnow)
        
        def __str__(self):
            return self.title
    
        meta = {"collection": "transaction_collection"}
    
    
    with MongoClient("mongodb://localhost:27017") as client:
        connection = client.get_database("transaction")
        collection = connection.transaction_collection
        transaction_set = TransactionQuerySet(Transaction, collection)
    
        print(transaction_set.order_by("-title"))
    

    OUTPUT

    [<Transaction: ETHUSTD>, <Transaction: BTCUSTD>, ...]