`Count` в наборе запросов неправильный после` extra`

Question

Когда я определенным образом использую extra в наборе запросов Django (назовите его qs), результат qs.count() отличается от len(qs.all()). Воспроизвести:

Создайте пустой проект и приложение Django, затем добавьте тривиальную модель:

class Baz(models.Model):
    pass

Теперь сделайте несколько объектов:

>>> Baz(id=1).save()
>>> Baz(id=2).save()
>>> Baz(id=3).save()
>>> Baz(id=4).save()

Использование метода extra для выбора только некоторых из них дает ожидаемое количество:

>>> Baz.objects.extra(where=['id > 2']).count()
2
>>> Baz.objects.extra(where=['-id < -2']).count()
2

Но добавьте предложение select к extra и обратитесь к нему в предложении where, и счет вдруг окажется неправильным, даже если результат all() верен:

>>> Baz.objects.extra(select={'negid': '0 - id'}, where=['"negid" < -2']).all()
[<Baz: Baz object>, <Baz: Baz object>]   # As expected
>>> Baz.objects.extra(select={'negid': '0 - id'}, where=['"negid" < -2']).count()
0   # Should be 2

Я думаю, что проблема связана с django.db.models.sql.query.BaseQuery.get_count(). Он проверяет, были ли установлены атрибуты BaseQuery select или aggregate_select; если это так, он использует подзапрос. Но django.db.models.sql.query.BaseQuery.add_extra добавляет только к атрибуту extra BaseQuery, а не select или aggregate_select.

Как я могу исправить проблему? Я знаю, что мог бы просто использовать len(qs.all()), но было бы неплохо иметь возможность передавать extra ed-набор запросов в другие части кода, и эти части могут вызывать count(), не зная, что он не работает.

Answer 1

Переопределение get_count() и обнаружение обезьян для устранения проблемы:

def get_count(self):
    """
    Performs a COUNT() query using the current filter constraints.
    """
    obj = self.clone()
    if len(self.select) > 1 or self.aggregate_select or self.extra:
        # If a select clause exists, then the query has already started to
        # specify the columns that are to be returned.
        # In this case, we need to use a subquery to evaluate the count.
        from django.db.models.sql.subqueries import AggregateQuery
        subquery = obj
        subquery.clear_ordering(True)
        subquery.clear_limits()

        obj = AggregateQuery(obj.model, obj.connection)
        obj.add_subquery(subquery)

    obj.add_count_column()
    number = obj.get_aggregation()[None]

    # Apply offset and limit constraints manually, since using LIMIT/OFFSET
    # in SQL (in variants that provide them) doesn't change the COUNT
    # output.
    number = max(0, number - self.low_mark)
    if self.high_mark is not None:
        number = min(number, self.high_mark - self.low_mark)

    return number

django.db.models.sql.query.BaseQuery.get_count = quuux.get_count

Тестирование:

>>> Baz.objects.extra(select={'negid': '0 - id'}, where=['"negid" < -2']).count()
2

Обновлено для работы с Django 1.2.1:

def basequery_get_count(self, using):
    """
    Performs a COUNT() query using the current filter constraints.
    """
    obj = self.clone()
    if len(self.select) > 1 or self.aggregate_select or self.extra:
        # If a select clause exists, then the query has already started to
        # specify the columns that are to be returned.
        # In this case, we need to use a subquery to evaluate the count.
        from django.db.models.sql.subqueries import AggregateQuery
        subquery = obj
        subquery.clear_ordering(True)
        subquery.clear_limits()

        obj = AggregateQuery(obj.model)
        obj.add_subquery(subquery, using=using)

    obj.add_count_column()
    number = obj.get_aggregation(using=using)[None]

    # Apply offset and limit constraints manually, since using LIMIT/OFFSET
    # in SQL (in variants that provide them) doesn't change the COUNT
    # output.
    number = max(0, number - self.low_mark)
    if self.high_mark is not None:
        number = min(number, self.high_mark - self.low_mark)

    return number
models.sql.query.Query.get_count = basequery_get_count

Однако я не уверен, что это исправление будет иметь другие непредвиденные последствия.