force_index and use_index PRIMARY clarify

Author: akvilibrium

src/django_mysql/models/query.py

@@ -262,6 +262,8 @@ def _index_hint(
 
         if len(index_names) == 0:
             indexes = "NONE"
+        elif len(index_names) == 1 and index_names[0] == 'PRIMARY':
+            indexes = 'PRIMARY'
         else:
             indexes = "`" + "`,`".join(index_names) + "`"
 

src/django_mysql/rewrite_query.py

@@ -33,7 +33,7 @@
         (?P<for_what>JOIN|ORDER\ BY|GROUP\ BY)
         \ # space
     )?
-    (?P<index_names>(`[^`]+`(,`[^`]+`)*)|NONE)
+    (?P<index_names>(`[^`]+`(,`[^`]+`)*)|NONE|PRIMARY)
     """,
     re.VERBOSE,
 )

tests/testapp/test_models.py

@@ -244,7 +244,7 @@ def test_use_index(self):
     def test_use_index_primary(self):
         with CaptureLastQuery() as cap:
             list(Author.objects.use_index("PRIMARY"))
-        assert ("USE INDEX (`PRIMARY`)") in cap.query
+        assert ("USE INDEX (PRIMARY)") in cap.query
         used = used_indexes(cap.query)
         assert len(used) == 0 or "PRIMARY" in used
 
@@ -273,10 +273,10 @@ def test_force_index(self):
 
     def test_force_index_primary(self):
         with CaptureLastQuery() as cap:
-            list(Author.objects.force_index("PRIMARY"))
-        assert ("FORCE INDEX (`PRIMARY`)") in cap.query
+            list(Author.objects.force_index("PRIMARY").order_by('pk'))
+        assert ("FORCE INDEX (PRIMARY)") in cap.query
         used = used_indexes(cap.query)
-        assert len(used) == 0 or "PRIMARY" in used
+        assert "PRIMARY" in used
 
     def test_force_index_inner_query(self):
         title_idx = index_name(Book, "title")
@@ -319,8 +319,8 @@ def test_ignore_index_multiple(self):
 
     def test_ignore_index_primary(self):
         with CaptureLastQuery() as cap:
-            list(Author.objects.filter(name__gt="").ignore_index("PRIMARY"))
-        assert "IGNORE INDEX (`PRIMARY`)" in cap.query
+            list(Author.objects.filter(name__gt="").ignore_index("PRIMARY").order_by('pk'))
+        assert "IGNORE INDEX (PRIMARY)" in cap.query
         assert "PRIMARY" not in used_indexes(cap.query)
 
     def test_force_index_at_least_one(self):
@@ -355,7 +355,7 @@ def test_use_index_table_name(self):
                     "PRIMARY", table_name=extra_table
                 )
             )
-        assert ("`" + extra_table + "` USE INDEX (`PRIMARY`) ") in cap.query
+        assert ("`" + extra_table + "` USE INDEX (PRIMARY) ") in cap.query
 
     def test_force_index_table_name_doesnt_exist_ignored(self):
         with CaptureLastQuery() as cap:

tests/testapp/test_rewrite_query.py

@@ -284,6 +284,15 @@ def test_index_hint_for_order_by(self):
             + "WHERE (1) ORDER BY col_a"
         )
 
+    def test_index_hint_primary(self):
+        assert rewrite_query(
+            "SELECT col_a, col_b FROM `sometable` "
+            + "WHERE (/*QueryRewrite':index=`sometable` FORCE PRIMARY */1)"
+        ) == (
+            "SELECT col_a, col_b FROM `sometable` "
+            + "FORCE INDEX (PRIMARY) WHERE (1)"
+        )
+
     def test_it_is_instrumented(self):
         with CaptureLastQuery() as cap, connection.cursor() as cursor:
             cursor.execute(