Search code examples
python-3.xmongoengineflask-mongoengine

Flask-Mongoengine custom order-by comparator


I am using Flask-Mongoengine as an ORM for my flask application. I define a Document to store unique versions in a string:

MyDocument(db.Document):
    version = db.StringField(primary_key=True, required=True)

The value stored in the StringField is in the following format:

a.b.c.d

I run into issues when it comes to using order_by for my queries:

MyDocument.objects.order_by('version')

More precisely, if the first part of the version (a in the format example above) contains multiple digits, it doesn't sort properly (e.g: 15.20.142 < 1.7.9). While I could obviously do something like:

tmp = Document.objects.all()
result = some_merge_sort(tmp)

However, that requires writing, testing and maintaining a sorting algorithm. This then brings me to my next solution, which is to overwrite the comparator used in order_by. The issue is, I'm not sure how to do this properly. Do I have to define a custom attribute type, or is there some hook I'm supposed to override?

In any case, I thought I would as the good folks on stackoverflow before I go into re-writing the wheel.


Solution

  • After doing some digging into the Mongoengine source code, I found out that it uses the PyMongo cursor sort which itself wraps the queries made to MongoDB. In other words, my Document definition is too high in the chain to even affect the order_by result.

    As such, the solution I went with was to write a classmethod. This simply takes in a mongoengine.queryset.QuerySet instance and applies a sort on the version attribute.

    from mongoengine.queryset import QuerySet
    from packaging import version
    
    class MyDocument(db.Document):
        version = db.StringField(primary_key=True, required=True)
    
        @classmethod
        def order_by_version(_, queryset, reverse=False):
            """
                Wrapper function to order a given queryset of this classes
                version attribute.
    
                :params:
                    - `QuerySet` :queryset: - The queryset to order the elements.
                    - `bool`     :reverse:  - If the results should be reversed.
    
                :raises:
                    - `TypeError` if the instance of the given queryset is not a `QuerySet`.
    
                :returns:
                    - `List` containing the elements sorted by the version attribute.
           """
           # Instance check the queryset.
           if not isinstance(queryset, QuerySet):
               raise TypeError("order_by_version requires a QuerySet instance!")
    
           # Sort by the version attribute.
           return sorted(queryset, key=lambda obj: version.parse(obj.version), reverse=reverse)
    

    Then to utilize this method, I can simply do as follows:

    MyDocument.order_by_version(MyDocument.objects)
    

    As for my reason for taking the QuerySet instance rather than doing the query from within the classmethod, it was simply to allow the opportunity to call other Mongoengine methods prior to doing the ordering. Because I used sorted, it will maintain any existing order, so I can still use the built-in order_by for other attributes.