Mongoengine custom queryset

5.1k 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 best solutions below

0
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
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