Mongoengine custom queryset

5k views Asked by At

I am trying to convert the ObjectId and ISODate to string representations while querying from MongoDB database.

def mongo_to_dict(obj, exclude_fields):
    return_data = []

    if obj is None:
        return None

    if isinstance(obj, Document):
        return_data.append(("_id", str(obj.id)))

    for field_name in obj._fields:

        if field_name in exclude_fields:
            continue

        if field_name in ("id",):
            continue

        data = obj._data[field_name]

        if isinstance(obj._fields[field_name], ListField):
            return_data.append((field_name, list_field_to_dict(data)))
        elif isinstance(obj._fields[field_name], EmbeddedDocumentField):
            return_data.append((field_name, mongo_to_dict(data, [])))
        elif isinstance(obj._fields[field_name], DictField):
            return_data.append((field_name, data))
        else:
            return_data.append((field_name, mongo_to_python_type(obj._fields[field_name], data)))

    return dict(return_data)


def list_field_to_dict(list_field):
    return_data = []

    for item in list_field:
        if isinstance(item, EmbeddedDocument):
            return_data.append(mongo_to_dict(item, []))
        else:
            return_data.append(mongo_to_python_type(item, item))

    return return_data


def mongo_to_python_type(field, data):
    if isinstance(field, DateTimeField):
        return time.mktime(data.timetuple()) * 1000
    elif isinstance(field, ComplexDateTimeField):
        return field.to_python(data).isoformat()
    elif isinstance(field, StringField):
        return str(data)
    elif isinstance(field, FloatField):
        return float(data)
    elif isinstance(field, IntField):
        return int(data)
    elif isinstance(field, BooleanField):
        return bool(data)
    elif isinstance(field, ObjectIdField):
        return str(data)
    elif isinstance(field, DecimalField):
        return data
    else:
        return str(data)


class Portfolio(Document):
    meta = {'collection': 'Portfolios'}
    PortfolioName = StringField()
    LastUpdateDate = DateTimeField(default=datetime.datetime.now())
    RecentActivity = ListField(default=[])


    def to_dict(self):
        return mongo_to_dict(self, [])

Now when i create a Portfolio Object like this

a = Portfolio(PortfolioName='BB Visa').save()

and when i try to get the to_dict() repressentation for the object like a.to_dict(), it works perfectly.

{'_id': '5581cf9129e457241a32e8f7', 'PortfolioName': 'BB Visa', 'RecentActivity': [], 'LastUpdateDate': 1434570641000.0}

But the problem is i want to_dict() to operate on class level not object level.

So when i tried to define a custom queryset like this

class Portfolio(Document):
    meta = {'collection': 'Portfolios'}
    PortfolioName = StringField()
    LastUpdateDate = DateTimeField(default=datetime.datetime.now())
    RecentActivity = ListField(default=[])

    @queryset_manager
    def to_dict(self, queryset):
        return mongo_to_dict(self, [])

Now running the following command Portfolio.to_dict() generates error like below Traceback (most recent call last):

  File "/home/ajay/PycharmProjects/solveit/test.py", line 102, in <module>
    print(Portfolio.to_dict())
  File "/home/ajay/.pyenv/versions/3.4.3/lib/python3.4/site-packages/mongoengine/queryset/manager.py", line 43, in __get__
    queryset = self.get_queryset(owner, queryset)
  File "/home/ajay/PycharmProjects/solveit/test.py", line 83, in to_dict
    return mongo_to_dict(self, [])
  File "/home/ajay/PycharmProjects/solveit/test.py", line 28, in mongo_to_dict
    data = obj._data[field_name]
TypeError: 'member_descriptor' object is not subscriptable

I understand that QuerySet Manager passes class into the function, thats why the error. How to solve this.

2

There are 2 answers

0
bellum On

Actually I am not sure that it is right decision to define this as @queryset_manager cause according to documentation

mongoengine.queryset.queryset_manager(func) Decorator that allows you to define custom QuerySet managers on Document classes. The manager must be a function that accepts a Document class as its first argument, and a QuerySet as its second argument. The method function should return a QuerySet, probably the same one that was passed in, but modified in some way.

I suggest adding some separate utils function that takes queryset:

def convert_queryset_to_list_of_dicts(queryset):
    return [mongo_to_dict(obj) for obj in queryset]

Or if you still want to have this as @queryset_manager then:

...
@queryset_manager
def to_dict(self, queryset):
    return [mongo_to_dict(obj) for obj in queryset]
...

Also there is a ready method of mongoengine Model objects to_mongo. You can try it like this: [obj.to_mongo() for obj in queryset]. Is not it applicable?

0
egvo On

That's an old question but it's very interesting and here still was no good answer. So here is my answer.

Intro

You're right in customizing queryset, but you shouldn't define another queryset_manager for your purposes. As @bellum mentioned @queryset_manager should return QuerySet. Instead you should define method inside QuerySet. You can do this inside your queryset_manager but it still won't work with filter, order_by, skip, limit and such. Instead you should define custom QuerSet and set it as default for model.

Answer

from mongoengine.queryset import QuerySet


class CustomQuerySet(QuerySet):

    # define your custom method here
    def print_ids(self):
        for doc in self:
            print(doc.id)


class Product(DynamicDocument):
    # set your CustomQuerySet as default queryset_class
    meta = {'queryset_class': CustomQuerySet}

Actually your model will use base QuerySet but now it will contain your custom methods.

Examples

>>> Product.objects.print_ids()
5e46433c943478954eef837d
5e46623705f1fa6c78e3fdbc
5e46627b05f1fa6c78e403a3

You may use it with filters:

>>> Product.objects.filter(id='5e46433c943478954eef837d').print_ids()
5e46433c943478954eef837d

You may use ordering, skip, limit and other stuff:

>>> Product.objects.all().limit(2).print_ids()
5e46433c943478954eef837d
5e46623705f1fa6c78e3fdbc