django get_or_create取决于自定义管理器方法的结果

Question

I am trying to implement a safe method of creating instances of the model Bar . I have a custom manager method to search for any nearby existing Bar instances based on a x, y pair of coordinates: Bar.objects.circle_search(x, y, radius)

These x, y values typically come from Foo . Upon creating a Foo instance it should then be associated via a foreign key with an existing Bar instance if Foo.x, Foo.y are within radius of Bar.x, Bar.y . If no Bar instances are found then create one. Each Bar instance has x, y columns that are defined by the Foo that created it.

I had something simple like:

bar = Bar.objects.circle_search(foo.x, foo.y, radius)
if bar is None:
    # fill some kwargs
    bar, created = Bar.objects.get_or_create(**kwargs)

however there is an obvious race here where two nearby foo 's could create separate Bar instances that are close to each other (there shouldn't be any Bar instances within radius of each other).

Is there anyway to chain/insert the circle_search method to be used as the get part of get_or_create ? ie in plain English - get a Bar that matches the circle_search , if not then create one.

# something like?
Bar.object.get_or_create(circle_search(foo.x, foo.y, radius), **kwargs)

The problem is that I cannot rely on using bar.x, bar.y (populated from foo.x, foo.y ) as unique since very nearby (in x, y space) Foo instances would be allowed to create Bar instances, but they shouldn't because I need to collect all Foo 's within radius to be associated with the same Bar .

(Part of me wonders if all this would be safe even with get_or_create since no IntegrityError would be raised even if we do end up making a Bar within radius of another - in that case is it possible to write some custom logic somewhere that will cause this to throw on any instance creation?)

Sorry if this is difficult to follow, I can try clarify in comments.

---

Minimal, simple example models for completeness:

class BarQuerySet(models.QuerySet):
    def circle_search(self, x, y, radius):
        # custom SQL to filter on a distance
        return qs

class Bar(models.Model):
    x = models.FloatField()
    y = models.FloatField()
    objects = BarQuerySet.as_manager()

class Foo(models.Model):
    x = models.FloatField()
    y = models.FloatField()
    bar = models.ForeignKey(Bar)

Answer 1

Your queries won't raise IntegrityError so you can't use get_or_create() , which relies on that. You could write a Constraint so that an IntegrityError is raised when another row within your constraint exists. But your lookup is probably too complex to pass to get_or_create() , in which case you should look at the Django codebase for get_or_create() and replicate it yourself with your own lookup.

Answer 2

I ended up adding some check in Bar.clean() and overriding the behaviour get_or_create in BarQuerySet , specifically _create_object_from_params .

Seems to return expected results, although not fully tested in a race situation but as it relies almost fully on django's implementation I guess that's safer than me cooking my own method.

class BarQuerySet(models.QuerySet):
    # ...
    def _create_object_from_params(self, lookup, params, lock=False):
        try:
            with transaction.atomic(using=self.db):
                params = {k: v() if callable(v) else v for k, v in params.items()}
                obj = self.create(**params)
            return obj, True
        except IntegrityError as e:
            try:
                qs = self.select_for_update() if lock else self
                # the default action is to try to get an existing instance from
                # the lookup kwargs (but this won't necessarily exist since x and 
                # y can be slightly different):
                # return qs.get(**lookup), False

                # instead, we now just grab the nearest source, continuing to
                # the raised IntegrityError if we don't grab it for any reason:
                obj = qs.circle_search(lookup['x'], lookup['y']).first()
                if obj is not None:
                    return obj, False
            except self.model.DoesNotExist:
                pass
            raise e


class Bar(models.Model):
    # ...
    def clean(self):
        # we only care about checking if we're creating 
        # a source, hence self._state.adding
        if self._state.adding and Bar.objects.circle_search(self.x, self.y, 1).exists():
            raise IntegrityError("existing nearby source found")
        super().clean()
    # also ensure Bar.save() calls `Bar.full_clean()` to perform this check always upon save

(This could perhaps have even been moved one up, to override Bar.objects.get() , if x, y are in the arguments to the filter.)