Search code examples
pythonflaskmongoengine

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


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)

Solution

  • 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.