How to let a Reference Field accept multiple Document schemas in MongoEngine?

1.7k views Asked by At

Context: I am writing an API (using Flask and MongoEngine) with multiple account types, including perhaps buildings. I need the database to hold some temporary accounts until a particular building registers.

This is how I've been referencing just one type of user:

current_holder_of_stuff = ReferenceField(ActiveUser)

I know GenericReferenceField is also an option, but what if I only want to allow two types of references? Is there anything like:

current_holder_of_stuff = ReferenceField(ActiveUser, TempUser)

Muchos thankos!

1

There are 1 answers

3
Steve Rossiter On BEST ANSWER

It may work to create a parent class of type User and then have inherited classes of ActiveUser and TempUser to deal with the various user types. As for the requirement for current_holder_of_stuff to be two possible document types, you cannot use a single reference field. As you've dismissed using GenericReferenceField then one way might be to add a property method and a StringField with options such as this:

import mongoegine as mdb

class User(mdb.Document):
    name = mdb.StringField()

    meta = {'allow_inheritance': True}

class ActiveUser(User):
    activation_date = mdb.DateTimeField()

class TempUser(User):
    date_limit = mdb.DateTimeField()

class Building(mdb.Document):
    address = mdb.StringField()

class Stuff(mdb.Document):
    user = mdb.ReferenceField(User)
    building = mdb.ReferenceField(Building)
    currently_with = mdb.StringField(options=['user','building'],required=True)

    @property
    def current_holder_of_stuff(self):
        if self.currently_with == "user":
            return self.user
        else:
            return self.building    

You can also use mongoengine's signals to perform checks pre-save to ensure there is only a user or building defined.