Ticket #16937: prefetch_3.1.diff

File prefetch_3.1.diff, 37.0 KB (added by Luke Plant, 13 years ago)

Fix to get tests to run when you run the whole test suite

  • django/contrib/contenttypes/generic.py

    diff --git a/django/contrib/contenttypes/generic.py b/django/contrib/contenttypes/generic.py
    a b  
    225225            content_type = content_type,
    226226            content_type_field_name = self.field.content_type_field_name,
    227227            object_id_field_name = self.field.object_id_field_name,
    228             core_filters = {
    229                 '%s__pk' % self.field.content_type_field_name: content_type.id,
    230                 '%s__exact' % self.field.object_id_field_name: instance._get_pk_val(),
    231             }
    232 
     228            prefetch_cache_name = self.field.attname,
    233229        )
    234230
    235231        return manager
     
    250246    """
    251247
    252248    class GenericRelatedObjectManager(superclass):
    253         def __init__(self, model=None, core_filters=None, instance=None, symmetrical=None,
     249        def __init__(self, model=None, instance=None, symmetrical=None,
    254250                     source_col_name=None, target_col_name=None, content_type=None,
    255                      content_type_field_name=None, object_id_field_name=None):
     251                     content_type_field_name=None, object_id_field_name=None,
     252                     prefetch_cache_name=None):
    256253
    257254            super(GenericRelatedObjectManager, self).__init__()
    258             self.core_filters = core_filters
    259255            self.model = model
    260256            self.content_type = content_type
    261257            self.symmetrical = symmetrical
     
    264260            self.target_col_name = target_col_name
    265261            self.content_type_field_name = content_type_field_name
    266262            self.object_id_field_name = object_id_field_name
     263            self.prefetch_cache_name = prefetch_cache_name
    267264            self.pk_val = self.instance._get_pk_val()
     265            self.core_filters = {
     266                '%s__pk' % content_type_field_name: content_type.id,
     267                '%s__exact' % object_id_field_name: instance._get_pk_val(),
     268            }
    268269
    269270        def get_query_set(self):
    270271            db = self._db or router.db_for_read(self.model, instance=self.instance)
    271272            return super(GenericRelatedObjectManager, self).get_query_set().using(db).filter(**self.core_filters)
    272273
     274        def get_prefetch_query_set(self, instances):
     275            if not instances:
     276                return self.model._default_manager.none()
     277
     278            db = self._db or router.db_for_read(self.model, instance=instances[0])
     279            query = {
     280                '%s__pk' % self.content_type_field_name: self.content_type.id,
     281                '%s__in' % self.object_id_field_name:
     282                    [obj._get_pk_val() for obj in instances]
     283                }
     284            return super(GenericRelatedObjectManager, self).get_query_set().using(db).filter(**query)
     285
     286        def select_matching_instances(self, obj, related_objects):
     287            pk_val = obj._get_pk_val()
     288            return [rel_obj for rel_obj in related_objects
     289                    if getattr(rel_obj, self.object_id_field_name) == pk_val]
     290
     291        def all(self):
     292            try:
     293                return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
     294            except (AttributeError, KeyError):
     295                return super(GenericRelatedObjectManager, self).all()
     296
     297
    273298        def add(self, *objs):
    274299            for obj in objs:
    275300                if not isinstance(obj, self.model):
  • django/db/models/fields/related.py

    diff --git a/django/db/models/fields/related.py b/django/db/models/fields/related.py
    a b  
    435435                db = self._db or router.db_for_read(self.model, instance=self.instance)
    436436                return super(RelatedManager, self).get_query_set().using(db).filter(**(self.core_filters))
    437437
     438            def get_prefetch_query_set(self, instances):
     439                """
     440                Return a queryset that does the bulk lookup needed
     441                by prefetch_related functionality.
     442                """
     443                if not instances:
     444                    return self.model._default_manager.none()
     445
     446                db = self._db or router.db_for_read(self.model, instance=instances[0])
     447                query = {'%s__%s__in' % (rel_field.name, attname):
     448                             [getattr(obj, attname) for obj in instances]}
     449                return super(RelatedManager, self).get_query_set().using(db).filter(**query)
     450
     451            def select_matching_instances(self, obj, related_objects):
     452                field_val = getattr(obj, attname)
     453                other_attname = rel_field.get_attname()
     454                return [rel_obj for rel_obj in related_objects
     455                        if getattr(rel_obj, other_attname) == field_val]
     456
     457            def all(self):
     458                try:
     459                    return self.instance._prefetched_objects_cache[rel_field.related_query_name()]
     460                except (AttributeError, KeyError):
     461                    return super(RelatedManager, self).all()
     462
    438463            def add(self, *objs):
    439464                for obj in objs:
    440465                    if not isinstance(obj, self.model):
     
    482507    """Creates a manager that subclasses 'superclass' (which is a Manager)
    483508    and adds behavior for many-to-many related objects."""
    484509    class ManyRelatedManager(superclass):
    485         def __init__(self, model=None, core_filters=None, instance=None, symmetrical=None,
     510        def __init__(self, model=None, query_field_name=None, instance=None, symmetrical=None,
    486511                     source_field_name=None, target_field_name=None, reverse=False,
    487                      through=None):
     512                     through=None, prefetch_cache_name=None):
    488513            super(ManyRelatedManager, self).__init__()
    489514            self.model = model
    490             self.core_filters = core_filters
     515            self.query_field_name = query_field_name
     516            self.core_filters = {'%s__pk' % query_field_name: instance._get_pk_val()}
    491517            self.instance = instance
    492518            self.symmetrical = symmetrical
    493519            self.source_field_name = source_field_name
    494520            self.target_field_name = target_field_name
    495521            self.reverse = reverse
    496522            self.through = through
     523            self.prefetch_cache_name = prefetch_cache_name
    497524            self._pk_val = self.instance.pk
    498525            if self._pk_val is None:
    499526                raise ValueError("%r instance needs to have a primary key value before a many-to-many relationship can be used." % instance.__class__.__name__)
     
    502529            db = self._db or router.db_for_read(self.instance.__class__, instance=self.instance)
    503530            return super(ManyRelatedManager, self).get_query_set().using(db)._next_is_sticky().filter(**(self.core_filters))
    504531
     532        def get_prefetch_query_set(self, instances):
     533            if not instances:
     534                return self.model._default_manager.none()
     535
     536            from django.db import connections
     537
     538            db = self._db or router.db_for_read(self.model, instance=instances[0])
     539            query = {'%s__pk__in' % self.query_field_name:
     540                         [obj._get_pk_val() for obj in instances]}
     541            qs = super(ManyRelatedManager, self).get_query_set().using(db)._next_is_sticky().filter(**query)
     542
     543            # M2M: need to annotate the query in order to get the PK of the
     544            # primary model that the secondary model was actually related to.
     545
     546            # We know that there will already be a join on the join table, so we
     547            # can just add the select.
     548            join_table = self.through._meta.db_table
     549            pk_col = "%s_id" % self.source_field_name
     550            connection = connections[db]
     551            qn = connection.ops.quote_name
     552            qs = qs.extra(select={'_prefetch_related_pk':
     553                                      '%s.%s' % (qn(join_table), qn(pk_col))})
     554            return qs
     555
     556        def select_matching_instances(self, obj, related_objects):
     557            pk_val = obj._get_pk_val()
     558            return [rel_obj for rel_obj in related_objects
     559                    if rel_obj._prefetch_related_pk == pk_val]
     560
     561        def all(self):
     562            try:
     563                return self.instance._prefetched_objects_cache[self.prefetch_cache_name]
     564            except (AttributeError, KeyError):
     565                return super(ManyRelatedManager, self).all()
     566
    505567        # If the ManyToMany relation has an intermediary model,
    506568        # the add and remove methods do not exist.
    507569        if rel.through._meta.auto_created:
     
    683745
    684746        manager = self.related_manager_cls(
    685747            model=rel_model,
    686             core_filters={'%s__pk' % self.related.field.name: instance._get_pk_val()},
     748            query_field_name=self.related.field.name,
     749            prefetch_cache_name=self.related.field.related_query_name(),
    687750            instance=instance,
    688751            symmetrical=False,
    689752            source_field_name=self.related.field.m2m_reverse_field_name(),
     
    739802
    740803        manager = self.related_manager_cls(
    741804            model=self.field.rel.to,
    742             core_filters={'%s__pk' % self.field.related_query_name(): instance._get_pk_val()},
     805            query_field_name=self.field.related_query_name(),
     806            prefetch_cache_name=self.field.name,
    743807            instance=instance,
    744808            symmetrical=self.field.rel.symmetrical,
    745809            source_field_name=self.field.m2m_field_name(),
  • django/db/models/manager.py

    diff --git a/django/db/models/manager.py b/django/db/models/manager.py
    a b  
    172172    def select_related(self, *args, **kwargs):
    173173        return self.get_query_set().select_related(*args, **kwargs)
    174174
     175    def prefetch_related(self, *args, **kwargs):
     176        return self.get_query_set().prefetch_related(*args, **kwargs)
     177
    175178    def values(self, *args, **kwargs):
    176179        return self.get_query_set().values(*args, **kwargs)
    177180
  • django/db/models/query.py

    diff --git a/django/db/models/query.py b/django/db/models/query.py
    a b  
    3636        self._iter = None
    3737        self._sticky_filter = False
    3838        self._for_write = False
     39        self._prefetch_related = set()
     40        self._prefetch_done = False
    3941
    4042    ########################
    4143    # PYTHON MAGIC METHODS #
     
    8183                self._result_cache = list(self.iterator())
    8284        elif self._iter:
    8385            self._result_cache.extend(self._iter)
     86        if self._prefetch_related and not self._prefetch_done:
     87            self._prefetch_related_objects()
    8488        return len(self._result_cache)
    8589
    8690    def __iter__(self):
     91        if self._prefetch_related:
     92            # We need all the results in order to be able to do the prefetch
     93            # in one go. To minimize code duplication, we use the __len__
     94            # code path which also forces this, and also does the prefetch
     95            len(self)
     96
    8797        if self._result_cache is None:
    8898            self._iter = self.iterator()
    8999            self._result_cache = []
     
    106116                self._fill_cache()
    107117
    108118    def __nonzero__(self):
     119        if self._prefetch_related:
     120            # We need all the results in order to be able to do the prefetch
     121            # in one go. To minimize code duplication, we use the __len__
     122            # code path which also forces this, and also does the prefetch
     123            len(self)
     124
    109125        if self._result_cache is not None:
    110126            return bool(self._result_cache)
    111127        try:
     
    526542            return self.query.has_results(using=self.db)
    527543        return bool(self._result_cache)
    528544
     545    def _prefetch_related_objects(self):
     546        # This method can only be called once the result cache has been filled.
     547        prefetch_related_objects(self._result_cache, self._prefetch_related)
     548        self._prefetch_done = True
     549
    529550    ##################################################
    530551    # PUBLIC METHODS THAT RETURN A QUERYSET SUBCLASS #
    531552    ##################################################
     
    649670            obj.query.max_depth = depth
    650671        return obj
    651672
     673    def prefetch_related(self, *fields):
     674        """
     675        Returns a new QuerySet instance that will prefetch Many-To-One
     676        and Many-To-Many related objects when the QuerySet is evaluated.
     677
     678        The fields specified must be attributes that return a RelatedManager of
     679        some kind when used on instances of the evaluated QuerySet.
     680
     681        These RelatedManagers will be modified so that their 'all()' method will
     682        return a QuerySet whose cache is already filled with objects that were
     683        looked up in a single batch, rather than one query per object in the
     684        current QuerySet.
     685
     686        When prefetch_related() is called more than once, the list of fields to
     687        prefetch is added to. If prefetch_related() is called with no arguments
     688        the list is cleared.
     689        """
     690        if fields:
     691            new_fields = self._prefetch_related.union(set(fields))
     692        else:
     693            new_fields = set()
     694        return self._clone(_prefetch_related=new_fields)
     695
    652696    def dup_select_related(self, other):
    653697        """
    654698        Copies the related selection status from the QuerySet 'other' to the
     
    798842            query.filter_is_sticky = True
    799843        c = klass(model=self.model, query=query, using=self._db)
    800844        c._for_write = self._for_write
     845        c._prefetch_related = self._prefetch_related
    801846        c.__dict__.update(kwargs)
    802847        if setup and hasattr(c, '_setup_query'):
    803848            c._setup_query()
     
    14841529    query = sql.InsertQuery(model)
    14851530    query.insert_values(fields, objs, raw=raw)
    14861531    return query.get_compiler(using=using).execute_sql(return_id)
     1532
     1533
     1534def prefetch_related_objects(result_cache, fields):
     1535    """
     1536    Populates prefetched objects caches for a list of results
     1537    from a QuerySet
     1538    """
     1539    from django.db.models.sql.constants import LOOKUP_SEP
     1540
     1541    if len(result_cache) == 0:
     1542        return # nothing to do
     1543
     1544    model = result_cache[0].__class__
     1545
     1546    # We need to be able to dynamically add to the list of prefetch_related
     1547    # fields that we look up (see below).  So we need some book keeping to
     1548    # ensure we don't do duplicate work.
     1549    done_fields = set() # list of fields like foo__bar__baz
     1550    done_lookups = {}   # dictionary of things like 'foo__bar': [results]
     1551    fields = list(fields)
     1552
     1553    # We may expand fields, so need a loop that allows for that
     1554    i = 0
     1555    while i < len(fields):
     1556        # 'field' can span several relationships, and so represent multiple
     1557        # lookups.
     1558        field = fields[i]
     1559
     1560        if field in done_fields:
     1561            # We've done exactly this already, skip the whole thing
     1562            i += 1
     1563            continue
     1564        done_fields.add(field)
     1565
     1566        # Top level, the list of objects to decorate is the the result cache
     1567        # from the primary QuerySet. It won't be for deeper levels.
     1568        obj_list = result_cache
     1569
     1570        attrs = field.split(LOOKUP_SEP)
     1571        for level, attr in enumerate(attrs):
     1572            # Prepare main instances
     1573            if len(obj_list) == 0:
     1574                break
     1575
     1576            good_objects = True
     1577            for obj in obj_list:
     1578                if not hasattr(obj, '_prefetched_objects_cache'):
     1579                    try:
     1580                        obj._prefetched_objects_cache = {}
     1581                    except AttributeError:
     1582                        # Must be in a QuerySet subclass that is not returning
     1583                        # Model instances, either in Django or 3rd
     1584                        # party. prefetch_related() doesn't make sense, so quit
     1585                        # now.
     1586                        good_objects = False
     1587                        break
     1588            if not good_objects:
     1589                break
     1590
     1591            # Descend down tree
     1592            try:
     1593                rel_obj = getattr(obj_list[0], attr)
     1594            except AttributeError:
     1595                raise AttributeError("Cannot find '%s' on %s object, '%s' is an invalid "
     1596                                     "parameter to prefetch_related()" %
     1597                                     (attr, obj_list[0].__class__.__name__, field))
     1598
     1599            can_prefetch = hasattr(rel_obj, 'get_prefetch_query_set')
     1600            if level == len(attrs) - 1 and not can_prefetch:
     1601                # Last one, this *must* resolve to a related manager.
     1602                raise ValueError("'%s' does not resolve to a supported 'many related"
     1603                                 " manager' for model %s - this is an invalid"
     1604                                 " parameter to prefetch_related()."
     1605                                 % (field, model.__name__))
     1606
     1607            if can_prefetch:
     1608                # Check we didn't do this already
     1609                lookup = LOOKUP_SEP.join(attrs[0:level+1])
     1610                if lookup in done_lookups:
     1611                    obj_list = done_lookups[lookup]
     1612                else:
     1613                    relmanager = rel_obj
     1614                    obj_list, additional_prf = _prefetch_one_level(obj_list, relmanager, attr)
     1615                    for f in additional_prf:
     1616                        new_prf = LOOKUP_SEP.join([lookup, f])
     1617                        fields.append(new_prf)
     1618                    done_lookups[lookup] = obj_list
     1619            else:
     1620                # Assume we've got some singly related object. We replace
     1621                # the current list of parent objects with that list.
     1622                obj_list = [getattr(obj, attr) for obj in obj_list]
     1623
     1624        i += 1
     1625
     1626
     1627def _prefetch_one_level(instances, relmanager, attname):
     1628    """
     1629    Runs prefetches on all instances using the manager relmanager,
     1630    assigning results to queryset against instance.attname.
     1631
     1632    The prefetched objects are returned, along with any additional
     1633    prefetches that must be done due to prefetch_related fields
     1634    found from default managers.
     1635    """
     1636    mainqs = relmanager.get_prefetch_query_set(instances)
     1637    # We have to handle the possibility that the default manager itself added
     1638    # prefetch_related fields to the QuerySet we just got back. We don't want to
     1639    # trigger the prefetch_related functionality by evaluating the query.
     1640    # Rather, we need to merge in the prefetch_related fields.
     1641    additional_prf = list(getattr(mainqs, '_prefetch_related', []))
     1642    if additional_prf:
     1643        mainqs = mainqs.prefetch_related()
     1644    all_related_objects = list(mainqs)
     1645    for obj in instances:
     1646        qs = getattr(obj, attname).all()
     1647        qs._result_cache = relmanager.select_matching_instances(obj, all_related_objects)
     1648        # We don't want the individual qs doing prefetch_related now, since we
     1649        # have merged this into the current work.
     1650        qs._prefetch_done = True
     1651        obj._prefetched_objects_cache[attname] = qs
     1652    return all_related_objects, additional_prf
  • docs/ref/models/querysets.txt

    diff --git a/docs/ref/models/querysets.txt b/docs/ref/models/querysets.txt
    a b  
    690690A :class:`~django.db.models.OneToOneField` is not traversed in the reverse
    691691direction if you are performing a depth-based ``select_related()`` call.
    692692
     693prefetch_related
     694~~~~~~~~~~~~~~~~
     695
     696.. method:: prefetch_related(*fields)
     697
     698.. versionadded:: 1.4
     699
     700Returns a ``QuerySet`` that will automatically retrieve, in a single batch,
     701related many-to-many and many-to-one objects for the specified fields.
     702
     703This is similar to ``select_related`` for the 'many related objects' case, but
     704note that ``prefetch_related`` causes a separate query to be issued for each set
     705of related objects that you request, unlike ``select_related`` which modifies
     706the original query with joins in order to get the related objects. With
     707``prefetch_related``, the additional queries are done as soon as the QuerySet
     708begins to be evaluated.
     709
     710For example, suppose you have these models::
     711
     712    class Topping(models.Model):
     713        name = models.CharField(max_length=30)
     714
     715    class Pizza(models.Model):
     716        name = models.CharField(max_length=50)
     717        toppings = models.ManyToManyField(Topping)
     718
     719        def __unicode__(self):
     720            return u"%s (%s)" % (self.name, u", ".join([topping.name
     721                                                        for topping in self.toppings.all()]))
     722
     723and run this code::
     724
     725    >>> Pizza.objects.all()
     726    [u"Hawaiian (ham, pineapple)", u"Seafood (prawns, smoked salmon)"...
     727
     728The problem with this code is that it will run a query on the Toppings table for
     729**every** item in the Pizza ``QuerySet``.  Using ``prefetch_related``, this can
     730be reduced to two:
     731
     732    >>> pizzas = Pizza.objects.all().prefetch_related('toppings')
     733
     734All the relevant toppings will be fetched in a single query, and used to make
     735``QuerySets`` that have a pre-filled cache of the relevant results. These
     736``QuerySets`` are then used in the ``self.toppings.all()`` calls.
     737
     738Please note that use of ``prefetch_related`` will mean that the additional
     739queries run will **always** be executed - even if you never use the related
     740objects - and it always fully populates the result cache on the primary
     741``QuerySet`` (which can sometimes be avoided in other cases).
     742
     743Also remember that, as always with QuerySets, any subsequent chained methods
     744will ignore previously cached results, and retrieve data using a fresh database
     745query. So, if you write the following:
     746
     747    >>> pizzas = Pizza.objects.prefetch_related('toppings')
     748    >>> [list(pizza.topppings.filter(spicy=True) for pizza in pizzas]
     749
     750...then the fact that `pizza.toppings.all()` has been prefetched will not help
     751you - in fact it hurts performance, since you have done a database query that
     752you haven't used. So use this feature with caution!
     753
     754The fields that must be supplied to this method can be any attributes on the
     755model instances which represent related queries that return multiple
     756objects. This includes attributes representing the 'many' side of ``ForeignKey``
     757relationships, forward and reverse ``ManyToManyField`` attributes, and also any
     758``GenericRelations``.
     759
     760You can also use the normal join syntax to do related fields of related
     761fields. Suppose we have an additional model to the example above::
     762
     763    class Restaurant(models.Model):
     764        pizzas = models.ManyToMany(Pizza, related_name='restaurants')
     765        best_pizza = models.ForeignKey(Pizza, related_name='championed_by')
     766
     767The following are all legal:
     768
     769    >>> Restaurant.objects.prefetch_related('pizzas__toppings')
     770
     771This will prefetch all pizzas belonging to restaurants, and all toppings
     772belonging to those pizzas. This will result in a total of 3 database queries -
     773one for the restaurants, one for the pizzas, and one for the toppings.
     774
     775    >>> Restaurant.objects.select_related('best_pizza').prefetch_related('best_pizza__toppings')
     776
     777This will fetch the best pizza and all the toppings for the best pizza for each
     778restaurant. This will be done in 2 database queries - one for the restaurants
     779and 'best pizzas' combined (achieved through use of ``select_related``), and one
     780for the toppings.
     781
     782Chaining ``prefetch_related`` calls will accumulate the fields that should have
     783this behavior applied. To clear any ``prefetch_related`` behavior, call the
     784method with no arguments.
     785
     786When using ``prefetch_related``, one difference is that, in some circumstances,
     787objects created by a query can be shared between the different objects that they
     788are related to i.e. a single Python model instance can appear at more than one
     789point in the tree of objects that are returned. Normally this behavior will not
     790be a problem, and will in fact save both memory and CPU time.
     791
    693792extra
    694793~~~~~
    695794
  • new file tests/modeltests/prefetch_related/models.py

    diff --git a/tests/modeltests/prefetch_related/__init__.py b/tests/modeltests/prefetch_related/__init__.py
    new file mode 100644
    diff --git a/tests/modeltests/prefetch_related/models.py b/tests/modeltests/prefetch_related/models.py
    new file mode 100644
    - +  
     1from django.contrib.contenttypes.models import ContentType
     2from django.contrib.contenttypes import generic
     3from django.db import models
     4
     5## Basic tests
     6
     7class Author(models.Model):
     8    name = models.CharField(max_length=50)
     9    first_book = models.ForeignKey('Book', related_name='first_time_authors')
     10
     11    def __unicode__(self):
     12        return self.name
     13
     14    class Meta:
     15        ordering = ['id']
     16
     17
     18class Book(models.Model):
     19    title = models.CharField(max_length=255)
     20    authors = models.ManyToManyField(Author, related_name='books')
     21
     22    def __unicode__(self):
     23        return self.title
     24
     25    class Meta:
     26        ordering = ['id']
     27
     28
     29class Reader(models.Model):
     30    name = models.CharField(max_length=50)
     31    books_read = models.ManyToManyField(Book, related_name='read_by')
     32
     33    def __unicode__(self):
     34        return self.name
     35
     36    class Meta:
     37        ordering = ['id']
     38
     39
     40## Models for default manager tests
     41
     42class Qualification(models.Model):
     43    name = models.CharField(max_length=10)
     44
     45    class Meta:
     46        ordering = ['id']
     47
     48
     49class TeacherManager(models.Manager):
     50    def get_query_set(self):
     51        return super(TeacherManager, self).get_query_set().prefetch_related('qualifications')
     52
     53
     54class Teacher(models.Model):
     55    name = models.CharField(max_length=50)
     56    qualifications = models.ManyToManyField(Qualification)
     57
     58    objects = TeacherManager()
     59
     60    def __unicode__(self):
     61        return "%s (%s)" % (self.name, ", ".join(q.name for q in self.qualifications.all()))
     62
     63    class Meta:
     64        ordering = ['id']
     65
     66
     67class Department(models.Model):
     68    name = models.CharField(max_length=50)
     69    teachers = models.ManyToManyField(Teacher)
     70
     71    class Meta:
     72        ordering = ['id']
     73
     74
     75## Generic relation tests
     76
     77class TaggedItem(models.Model):
     78    tag = models.SlugField()
     79    content_type = models.ForeignKey(ContentType, related_name="taggeditem_set2")
     80    object_id = models.PositiveIntegerField()
     81    content_object = generic.GenericForeignKey('content_type', 'object_id')
     82
     83    def __unicode__(self):
     84        return self.tag
     85
     86
     87class Bookmark(models.Model):
     88    url = models.URLField()
     89    tags = generic.GenericRelation(TaggedItem)
  • new file tests/modeltests/prefetch_related/tests.py

    diff --git a/tests/modeltests/prefetch_related/tests.py b/tests/modeltests/prefetch_related/tests.py
    new file mode 100644
    - +  
     1from django.contrib.contenttypes.models import ContentType
     2from django.test import TestCase
     3
     4from models import Author, Book, Reader, Qualification, Teacher, Department, TaggedItem, Bookmark
     5
     6
     7class PrefetchRelatedTests(TestCase):
     8
     9    def setUp(self):
     10
     11        self.book1 = Book.objects.create(title="Poems")
     12        self.book2 = Book.objects.create(title="Jane Eyre")
     13        self.book3 = Book.objects.create(title="Wuthering Heights")
     14        self.book4 = Book.objects.create(title="Sense and Sensibility")
     15
     16        self.author1 = Author.objects.create(name="Charlotte",
     17                                             first_book=self.book1)
     18        self.author2 = Author.objects.create(name="Anne",
     19                                             first_book=self.book1)
     20        self.author3 = Author.objects.create(name="Emily",
     21                                             first_book=self.book1)
     22        self.author4 = Author.objects.create(name="Jane",
     23                                             first_book=self.book4)
     24
     25        self.book1.authors.add(self.author1, self.author2, self.author3)
     26        self.book2.authors.add(self.author1)
     27        self.book3.authors.add(self.author3)
     28        self.book4.authors.add(self.author4)
     29
     30        self.reader1 = Reader.objects.create(name="Amy")
     31        self.reader2 = Reader.objects.create(name="Belinda")
     32
     33        self.reader1.books_read.add(self.book1, self.book4)
     34        self.reader2.books_read.add(self.book2, self.book4)
     35
     36    def test_m2m_forward(self):
     37        with self.assertNumQueries(2):
     38            lists = [list(b.authors.all()) for b in Book.objects.prefetch_related('authors')]
     39
     40        normal_lists = [list(b.authors.all()) for b in Book.objects.all()]
     41        self.assertEqual(lists, normal_lists)
     42
     43
     44    def test_m2m_reverse(self):
     45        with self.assertNumQueries(2):
     46            lists = [list(a.books.all()) for a in Author.objects.prefetch_related('books')]
     47
     48        normal_lists = [list(a.books.all()) for a in Author.objects.all()]
     49        self.assertEqual(lists, normal_lists)
     50
     51    def test_foreignkey_reverse(self):
     52        with self.assertNumQueries(2):
     53            lists = [list(b.first_time_authors.all())
     54                     for b in Book.objects.prefetch_related('first_time_authors')]
     55
     56        self.assertQuerysetEqual(self.book2.authors.all(), [u"<Author: Charlotte>"])
     57
     58    def test_survives_clone(self):
     59        with self.assertNumQueries(2):
     60            lists = [list(b.first_time_authors.all())
     61                     for b in Book.objects.prefetch_related('first_time_authors').exclude(id=1000)]
     62
     63    def test_len(self):
     64        with self.assertNumQueries(2):
     65            qs = Book.objects.prefetch_related('first_time_authors')
     66            length = len(qs)
     67            lists = [list(b.first_time_authors.all())
     68                     for b in qs]
     69
     70    def test_bool(self):
     71        with self.assertNumQueries(2):
     72            qs = Book.objects.prefetch_related('first_time_authors')
     73            x = bool(qs)
     74            lists = [list(b.first_time_authors.all())
     75                     for b in qs]
     76
     77    def test_clear(self):
     78        """
     79        Test that we can clear the behavior by calling prefetch_related()
     80        """
     81        with self.assertNumQueries(5):
     82            with_prefetch = Author.objects.prefetch_related('books')
     83            without_prefetch = with_prefetch.prefetch_related()
     84            lists = [list(a.books.all()) for a in without_prefetch]
     85
     86    def test_m2m_then_m2m(self):
     87        """
     88        Test we can follow a m2m and another m2m
     89        """
     90        with self.assertNumQueries(3):
     91            qs = Author.objects.prefetch_related('books__read_by')
     92            lists = [[[unicode(r) for r in b.read_by.all()]
     93                      for b in a.books.all()]
     94                     for a in qs]
     95            self.assertEqual(lists,
     96            [
     97                [[u"Amy"], [u"Belinda"]],  # Charlotte - Poems, Jane Eyre
     98                [[u"Amy"]],                # Anne - Poems
     99                [[u"Amy"], []],            # Emily - Poems, Wuthering Heights
     100                [[u"Amy", u"Belinda"]],    # Jane - Sense and Sense
     101            ])
     102
     103    def test_get(self):
     104        """
     105        Test that objects retrieved with .get() get the prefetch behaviour
     106        """
     107        # Need a double
     108        with self.assertNumQueries(3):
     109            author = Author.objects.prefetch_related('books__read_by').get(name="Charlotte")
     110            lists = [[unicode(r) for r in b.read_by.all()]
     111                      for b in author.books.all()]
     112            self.assertEqual(lists, [[u"Amy"], [u"Belinda"]])  # Poems, Jane Eyre
     113
     114    def test_foreign_key_then_m2m(self):
     115        """
     116        Test we can follow an m2m relation after a relation like ForeignKey
     117        that doesn't have many objects
     118        """
     119
     120        with self.assertNumQueries(2):
     121            qs = Author.objects.select_related('first_book').prefetch_related('first_book__read_by')
     122            lists = [[unicode(r) for r in a.first_book.read_by.all()]
     123                     for a in qs]
     124            self.assertEqual(lists, [[u"Amy"],
     125                                     [u"Amy"],
     126                                     [u"Amy"],
     127                                     [u"Amy", "Belinda"]])
     128
     129    def test_reuse(self):
     130        # Check re-use of objects.
     131        qs1 = Reader.objects.all()
     132        qs2 = Reader.objects.prefetch_related('books_read__first_time_authors')
     133
     134        authors1 = [a for r in qs1
     135                    for b in r.books_read.all()
     136                    for a in b.first_time_authors.all()]
     137        authors2 = [a for r in qs2
     138                    for b in r.books_read.all()
     139                    for a in b.first_time_authors.all()]
     140
     141        # The second prefetch_related lookup is a reverse foreign key. This
     142        # means the query for it can only be something like
     143        # "first_time_authors__pk__in = [...]" and cannot return more rows then
     144        # the number of Author objects in the database.  This means that these
     145        # objects will be reused (since in our data we've arranged for there
     146        # len(authors1) > Author.objects.count())
     147
     148        total_authors = Author.objects.count()
     149        self.assertEqual(len(authors1), len(authors2))
     150        self.assertTrue(len(authors1) > total_authors)
     151        self.assertTrue(len(set(map(id, authors1))) > len(set(map(id, authors2))))
     152        self.assertEqual(total_authors, len(set(map(id, authors2))))
     153
     154    def test_attribute_error(self):
     155        qs = Reader.objects.all().prefetch_related('books_read__xyz')
     156        with self.assertRaises(AttributeError) as cm:
     157            list(qs)
     158
     159        self.assertTrue('prefetch_related' in cm.exception.message)
     160
     161    def test_invalid_final_lookup(self):
     162        qs = Book.objects.prefetch_related('authors__first_book')
     163        with self.assertRaises(ValueError) as cm:
     164            list(qs)
     165
     166        self.assertTrue('prefetch_related' in cm.exception.message)
     167        self.assertTrue("first_book" in cm.exception.message)
     168
     169
     170class DefaultManagerTests(TestCase):
     171
     172    def setUp(self):
     173        self.qual1 = Qualification.objects.create(name="BA")
     174        self.qual2 = Qualification.objects.create(name="BSci")
     175        self.qual3 = Qualification.objects.create(name="MA")
     176        self.qual4 = Qualification.objects.create(name="PhD")
     177
     178        self.teacher1 = Teacher.objects.create(name="Mr Cleese")
     179        self.teacher2 = Teacher.objects.create(name="Mr Idle")
     180        self.teacher3 = Teacher.objects.create(name="Mr Chapman")
     181
     182        self.teacher1.qualifications.add(self.qual1, self.qual2, self.qual3, self.qual4)
     183        self.teacher2.qualifications.add(self.qual1)
     184        self.teacher3.qualifications.add(self.qual2)
     185
     186        self.dept1 = Department.objects.create(name="English")
     187        self.dept2 = Department.objects.create(name="Physics")
     188
     189        self.dept1.teachers.add(self.teacher1, self.teacher2)
     190        self.dept2.teachers.add(self.teacher1, self.teacher3)
     191
     192    def test_m2m_then_m2m(self):
     193        with self.assertNumQueries(3):
     194            # When we prefetch the teachers, and force the query, we don't want
     195            # the default manager on teachers to immediately get all the related
     196            # qualifications, since this will do one query per teacher.
     197            qs = Department.objects.prefetch_related('teachers')
     198            depts = "".join(["%s department: %s\n" %
     199                             (dept.name, ", ".join(unicode(t) for t in dept.teachers.all()))
     200                             for dept in qs])
     201
     202            self.assertEqual(depts,
     203                             "English department: Mr Cleese (BA, BSci, MA, PhD), Mr Idle (BA)\n"
     204                             "Physics department: Mr Cleese (BA, BSci, MA, PhD), Mr Chapman (BSci)\n")
     205
     206
     207class GenericRelationTests(TestCase):
     208
     209    def test_traverse_GFK(self):
     210        """
     211        Test that we can traverse a 'content_object' with prefetch_related()
     212        """
     213        # In fact, there is no special support for this in prefetch_related code
     214        # - we can traverse any object that will lead us to objects that have
     215        # related managers.
     216
     217        book1 = Book.objects.create(title="Winnie the Pooh")
     218        book2 = Book.objects.create(title="Do you like green eggs and spam?")
     219
     220        reader1 = Reader.objects.create(name="me")
     221        reader2 = Reader.objects.create(name="you")
     222
     223        book1.read_by.add(reader1)
     224        book2.read_by.add(reader2)
     225
     226        TaggedItem.objects.create(tag="awesome", content_object=book1)
     227        TaggedItem.objects.create(tag="awesome", content_object=book2)
     228
     229        ct = ContentType.objects.get_for_model(Book)
     230
     231        # We get 4 queries - 1 for main query, 2 for each access to
     232        # 'content_object' because these can't be handled by select_related, and
     233        # 1 for the 'read_by' relation.
     234        with self.assertNumQueries(4):
     235            # If we limit to books, we know that they will have 'read_by'
     236            # attributes, so the following makes sense:
     237            qs = TaggedItem.objects.select_related('content_type').prefetch_related('content_object__read_by').filter(tag='awesome').filter(content_type=ct, tag='awesome')
     238            readers_of_awesome_books = [r.name for tag in qs
     239                                        for r in tag.content_object.read_by.all()]
     240            self.assertEqual(readers_of_awesome_books, ["me", "you"])
     241
     242
     243    def test_generic_relation(self):
     244        b = Bookmark.objects.create(url='http://www.djangoproject.com/')
     245        t1 = TaggedItem.objects.create(content_object=b, tag='django')
     246        t2 = TaggedItem.objects.create(content_object=b, tag='python')
     247
     248        with self.assertNumQueries(2):
     249            tags = [t.tag for b in Bookmark.objects.prefetch_related('tags')
     250                    for t in b.tags.all()]
     251            self.assertEqual(sorted(tags), ["django", "python"])
Back to Top