Repository: kayak/pypika
Pull request #470 (2020-08-03 14:51:44)

Add use index option
Hello,

`USE INDEX` comes handy when you need to optimize your queries. I will be happy if you can check and merge this.

Thanks


Patch:
diff --git a/pypika/queries.py b/pypika/queries.py
index 8406967..0ef9b54 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -594,6 +594,7 @@ class QueryBuilder(Selectable, Term):
         self._with = []
         self._selects = []
         self._force_indexes = []
+        self._use_indexes = []
         self._columns = []
         self._values = []
         self._distinct = False
@@ -828,6 +829,14 @@ class QueryBuilder(Selectable, Term):
             elif isinstance(t, str):
                 self._force_indexes.append(Index(t))
 
+    @builder
+    def use_index(self, term: Union[str, Index], *terms: Union[str, Index]) -> "QueryBuilder":
+        for t in (term, *terms):
+            if isinstance(t, Index):
+                self._use_indexes.append(t)
+            elif isinstance(t, str):
+                self._use_indexes.append(Index(t))
+
     @builder
     def distinct(self) -> "QueryBuilder":
         self._distinct = True
@@ -1237,6 +1246,9 @@ class QueryBuilder(Selectable, Term):
         if self._force_indexes:
             querystring += self._force_index_sql(**kwargs)
 
+        if self._use_indexes:
+            querystring += self._use_index_sql(**kwargs)
+
         if self._joins:
             querystring += " " + " ".join(
                   join.get_sql(**kwargs) for join in self._joins
@@ -1361,6 +1373,11 @@ class QueryBuilder(Selectable, Term):
               indexes=",".join(index.get_sql(**kwargs) for index in self._force_indexes),
         )
 
+    def _use_index_sql(self, **kwargs: Any) -> str:
+        return " USE INDEX ({indexes})".format(
+              indexes=",".join(index.get_sql(**kwargs) for index in self._use_indexes),
+        )
+
     def _prewhere_sql(self, quote_char: Optional[str] = None, **kwargs: Any) -> str:
         return " PREWHERE {prewhere}".format(
               prewhere=self._prewheres.get_sql(


Test patch:
diff --git a/pypika/tests/test_selects.py b/pypika/tests/test_selects.py
index c629585..a463bc6 100644
--- a/pypika/tests/test_selects.py
+++ b/pypika/tests/test_selects.py
@@ -236,6 +236,27 @@ class SelectTests(unittest.TestCase):
 
         self.assertEqual('SELECT "foo" FROM "abc" FORCE INDEX ("egg","spam")', str(q))
 
+    def test_select_with_use_index(self):
+        q = Query.from_("abc").select("foo").use_index("egg")
+
+        self.assertEqual('SELECT "foo" FROM "abc" USE INDEX ("egg")', str(q))
+
+    def test_select_with_use_index_with_index_object(self):
+        index = Index("egg")
+        q = Query.from_("abc").select("foo").use_index(index)
+
+        self.assertEqual('SELECT "foo" FROM "abc" USE INDEX ("egg")', str(q))
+
+    def test_select_with_use_index_multiple_indexes(self):
+        q = Query.from_("abc").select("foo").use_index("egg", "bacon")
+
+        self.assertEqual('SELECT "foo" FROM "abc" USE INDEX ("egg","bacon")', str(q))
+
+    def test_select_with_use_index_multiple_calls(self):
+        q = Query.from_("abc").select("foo").use_index("egg", ).use_index("spam")
+
+        self.assertEqual('SELECT "foo" FROM "abc" USE INDEX ("egg","spam")', str(q))
+
     def test_mysql_query_uses_backtick_quote_chars(self):
         q = MySQLQuery.from_("abc").select("foo", "bar")