override Django get_or_create

Question

I have a model that I overrode the save method for, so that the save method can be passed in some data and auto-fill-in a field before saving. Here is my model:

class AccountModel(models.Model):

    account = models.ForeignKey(Account)

    def save(self, request=None, *args, **kwargs):
        if request:
            self.account = request.session['account']
        super(AccountModel, self).save(*args, **kwargs)

    class Meta:
        abstract = True

The idea is I set up a base model for objects that need to be associated with an account and then I won't have to deal with the account connections every time they come up (which is a lot).

But: I'd also like to use get_or_create, which saves the new objects without passing in the request. I know it is doable to not use get_or_create and do a try/except instead, but I'd like to know if there is a way to override get_or_create and what is the proper way to do it.

I looked at the code for the Manager (which I am looking at overriding) and the get_or_create function just calls a QuerySet.get_or_create function. Maybe I can write it to use other manager functions and not the QuerySet version of get_or_create? What do y'all think?

Answer 1

You could subclass django.db.models.query.QuerySet and override the get_or_create method there to accept your request keyword argument and pass it onto save I guess, but it isn't very pretty.

class AccountQuerySet(models.query.QuerySet):
    def get_or_create(...):
        ...

You could then add a custom manager to your Account model which uses this custom QuerySet:

class AccountManager(models.Manager):
    def get_queryset(self):
        return AccountQuerySet(self.model)

Then use this manager in your model:

class Account(models.Model):
    ...
    objects = AccountManager()

But you might find that the try-except method is neater after all :)