from django.db import models
# Some convenience functions.
def _get_table(model):
return model._meta.db_table
def _get_column(model, attr):
return "%s.%s" % (_get_table(model), model._meta.get_field(attr).column)
def _get_related_column(model, attr):
related = model._meta.get_field(attr).rel
return "%s.%s" % (_get_table(related.to), related.field_name)
def _related_count_sql(related, attr):
return """SELECT COUNT(*) FROM %s WHERE %s = %s
""" % (_get_table(related), _get_column(related, attr),
_get_related_column(related, attr))
# This is what it's all about.
def _count_related(self, query, count_attr=None, related_attr=None):
"""
Count the rows matching `query` related to this model by their
foreign key attribute `related_attr`, and store the result in
`count_attr`.
If `count_attr` is None, use the name of the module given in `query`
suffixed with '__count'.
If `related_attr` is None, find the first foreign key field in the
model queried by `query` relating to this model.
If `query` is a model class, use the all() method on its default
manager as the query.
"""
if isinstance(query, models.base.ModelBase):
query = query._default_manager.all()
if count_attr is None:
count_attr = query.model._meta.module_name + '__count'
if related_attr is None:
for field in query.model._meta.fields:
if isinstance(field, models.related.RelatedField):
if field.rel.to is self.model:
related_attr = field.name
select_count = _related_count_sql(query.model, related_attr)
joins, wheres, params = query._filters.get_sql(query.model._meta)
if wheres:
select_count += ' AND %s' % wheres[0]
return self.extra(select={count_attr: select_count}, params=params)
# Here's where your code comes in.
class YourManager(models.Manager):
count_related = _count_related # Add this to your managers.
Comments
Hmm, it seems that
filterdoesn't actually see the attribute added bycount_related. It must be possible sinceorder_bysees it just fine... probably a shortcoming ofQuerySet.#
I'm definitely interested in getting this to work with many-to-many relationships. What would be involved in doing so?
#