Repository: adamchainz/django-mysql
Issue: Query hints prevent approx_count from working

The way the ORM is wrapped to add query hints uses `extra()` to add where clauses; these are then seen by `approx_count` as preventing approximate counting, even though it could be done (if there are no other where clauses). This should just require an extension of the check whether the queryset has any filters to check only non-query hint filters.


Changed files:
tests/django_mysql_tests/test_models.py, django_mysql/models/query.py

diff --git a/django_mysql/models/query.py b/django_mysql/models/query.py
index 8447e61..d017942 100644
--- a/django_mysql/models/query.py
+++ b/django_mysql/models/query.py
@@ -6,6 +6,7 @@ from copy import copy
 from subprocess import PIPE, Popen
 
 from django.db import connections, models
+from django.db.models.sql.where import ExtraWhere
 from django.db.transaction import atomic
 from django.test.utils import CaptureQueriesContext
 from django.utils import six
@@ -13,6 +14,7 @@ from django.utils.functional import cached_property
 from django.utils.translation import ugettext as _
 
 from django_mysql.models.handler import Handler
+from django_mysql.rewrite_query import REWRITE_MARKER
 from django_mysql.status import GlobalStatus
 from django_mysql.utils import (
     StopWatch, WeightedAverageRate, have_program, noop_context,
@@ -363,32 +365,39 @@ class SmartPKRangeIterator(SmartChunkedIterator):
 def approx_count(queryset):
     # Returns the approximate count or raises a ValueError if this queryset
     # cannot be approximately counted
-    query = queryset.query
-
-    can_approx_count = (
-        not query.where and
-        query.high_mark is None and
-        query.low_mark == 0 and
-        not query.select and
-        not query.group_by and
-        not query.distinct
-    )
-
-    if hasattr(query, 'having'):  # Django < 1.9
-        can_approx_count = (can_approx_count and not query.having)
-
-    if not can_approx_count:
+    if not can_approx_count(queryset):
         raise ValueError("This QuerySet cannot be approximately counted")
 
     connection = connections[queryset.db]
     with connection.cursor() as cursor:
         table_name = queryset.model._meta.db_table
-        query = "EXPLAIN SELECT COUNT(*) FROM `{0}`".format(table_name)
-        cursor.execute(query)
+        sql = "EXPLAIN SELECT COUNT(*) FROM `{0}`".format(table_name)
+        cursor.execute(sql)
         approx_count = cursor.fetchone()[8]  # 'rows' is the 9th column
         return approx_count
 
 
+def can_approx_count(queryset):
+    query = queryset.query
+
+    if query.select or query.group_by or query.distinct:
+        return False
+    elif query.high_mark is not None or query.low_mark != 0:
+        return False
+
+    # Visit parts of the where clause - if any is not a query hint, fail
+    for child in query.where.children:
+        if not isinstance(child, ExtraWhere):
+            return False
+        elif not all((REWRITE_MARKER in sql) for sql in child.sqls):
+            return False
+
+    if hasattr(query, 'having') and query.having:  # Django < 1.9
+        return False  # pragma: no cover
+
+    return True
+
+
 def pt_visual_explain(queryset, display=True):
     if not have_program('pt-visual-explain'):  # pragma: no cover
         raise OSError("pt-visual-explain doesn't appear to be installed")
diff --git a/tests/django_mysql_tests/test_models.py b/tests/django_mysql_tests/test_models.py
index 79e8e32..5a7863d 100644
--- a/tests/django_mysql_tests/test_models.py
+++ b/tests/django_mysql_tests/test_models.py
@@ -84,6 +84,23 @@ class ApproximateCountTests(TransactionTestCase):
         with pytest.raises(ValueError):
             Author.objects.distinct().approx_count(fall_back=False)
 
+    def test_fallback_with_arbitrary_extra(self):
+        assert Author.objects.extra(where=['1=1']).approx_count() == 10
+        with pytest.raises(ValueError):
+            Author.objects.extra(where=['1=1']).approx_count(fall_back=False)
+
+    def test_approx_count_with_label(self):
+        # It should be possible to approx count a query set with a query hint
+        # as none of them affect the result
+        assert Author.objects.label('bla').approx_count(fall_back=False) == 10
+
+    def test_approx_count_with_straight_join(self):
+        # It should be possible to approx count a query set with a query hint
+        # as none of them affect the result
+        assert (
+            Author.objects.straight_join().approx_count(fall_back=False) == 10
+        )
+
 
 class QueryHintTests(TransactionTestCase):