coleifer/peewee — Issue #27: SelectQuery.count() problem
Author: awahlig
Created: 2011-12-16 01:39:51+00:00
State: open

I was doing a query using annotate() (with the default Count aggregation) to order by the count of related models.

The query itself worked well but calling query.count() failed with:

```
OperationalError: no such column: count
```

Further investigation revealed that SelectQuery.count() replaces the query with bare COUNT(x) but doesn't remove the ORDER BY clause which in my case created an invalid statement referencing an unknown column alias.

Here's how I fixed it:

```python

def count(self):
tmp = self.query, self._limit, self._offset, self._order_by
self._limit = self._offset = None
self._order_by = []
try:

if self.use_aliases():
self.query = 'COUNT(t1.%s)' % (self.model._meta.pk_name)
else:
self.query = 'COUNT(%s)' % (self.model._meta.pk_name)

res = self.database.execute(*self.sql())

finally:
# restore
self.query, self._limit, self._offset, self._order_by = tmp

return (res.fetchone() or [0])[0]
```
Two other things that I changed are:

* check whether fetchone() returned None (which is the case if the table is empty) and return 0 in that case
* use try..finally block to make sure the state is restored in case of an error
* replace multiple tmp variables with one