Repository: kayak/pypika
Pull request #334 (2019-10-17 14:59:41)

Add force index
Add FORCE INDEX hint to enhance performance with query optimizer

This replaces PR 'add force index hint #330' (I made the changes on branch in my fork - whereas #330 PR was done on patch, thus this new PR)

I added class for Index (previously, I used Field) and tests. Requests and comments from #330 were taken in account.


Patch:
diff --git a/.gitignore b/.gitignore
index 002c8b1..f7aa02f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -72,4 +72,4 @@ venv/
 ENV/
 
 # OS X
-.DS_Store
\ No newline at end of file
+.DS_Store
diff --git a/pypika/__init__.py b/pypika/__init__.py
index d404497..c5d37b5 100644
--- a/pypika/__init__.py
+++ b/pypika/__init__.py
@@ -65,6 +65,7 @@ from pypika.terms import (
     Criterion,
     EmptyCriterion,
     Field,
+    Index,
     Interval,
     JSON,
     Not,
diff --git a/pypika/queries.py b/pypika/queries.py
index 3b2eb06..7e7975e 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -10,6 +10,7 @@ from pypika.terms import (
     EmptyCriterion,
     Field,
     Function,
+    Index,
     Rollup,
     Star,
     Term,
@@ -440,6 +441,7 @@ class QueryBuilder(Selectable, Term):
 
         self._with = []
         self._selects = []
+        self._force_indexes = []
         self._columns = []
         self._values = []
         self._distinct = False
@@ -624,6 +626,14 @@ class QueryBuilder(Selectable, Term):
             self._validate_terms_and_append(*terms)
         self._replace = True
 
+    @builder
+    def force_index(self, term, *terms):
+        for t in (term, *terms):
+            if isinstance(t, Index):
+                self._force_indexes.append(t)
+            elif isinstance(t, str):
+                self._force_indexes.append(Index(t))
+
     @builder
     def distinct(self):
         self._distinct = True
@@ -954,6 +964,9 @@ class QueryBuilder(Selectable, Term):
         if self._from:
             querystring += self._from_sql(**kwargs)
 
+        if self._force_indexes:
+            querystring += self._force_index_sql(**kwargs)
+
         if self._joins:
             querystring += " " + " ".join(join.get_sql(**kwargs)
                                           for join in self._joins)
@@ -1055,6 +1068,12 @@ class QueryBuilder(Selectable, Term):
               for clause in self._from
         ))
 
+    def _force_index_sql(self, **kwargs):
+        return ' FORCE INDEX ({indexes})'.format(indexes=','.join(
+            index.get_sql(**kwargs)
+            for index in self._force_indexes),
+        )
+
     def _prewhere_sql(self, quote_char=None, **kwargs):
         return ' PREWHERE {prewhere}'.format(
               prewhere=self._prewheres.get_sql(quote_char=quote_char, subquery=True, **kwargs))
diff --git a/pypika/terms.py b/pypika/terms.py
index 97eaa56..b1b0061 100644
--- a/pypika/terms.py
+++ b/pypika/terms.py
@@ -478,6 +478,15 @@ class Field(Criterion, JSON):
         return field_sql
 
 
+class Index(Term):
+    def __init__(self, name, alias=None):
+        super(Index, self).__init__(alias)
+        self.name = name
+
+    def get_sql(self, quote_char=None, **kwargs):
+        return format_quotes(self.name, quote_char)
+
+
 class Star(Field):
     def __init__(self, table=None):
         super(Star, self).__init__('*', table=table)


Test patch:
diff --git a/pypika/tests/test_selects.py b/pypika/tests/test_selects.py
index c7dc5a8..687f731 100644
--- a/pypika/tests/test_selects.py
+++ b/pypika/tests/test_selects.py
@@ -6,6 +6,7 @@ from pypika import (
     ClickHouseQuery,
     EmptyCriterion,
     Field as F,
+    Index,
     MSSQLQuery,
     MySQLQuery,
     NullValue,
@@ -197,6 +198,27 @@ class SelectTests(unittest.TestCase):
 
         self.assertEqual('SELECT "foo" FROM "abc" LIMIT 10 OFFSET 10', str(q1))
 
+    def test_select_with_force_index(self):
+        q = Query.from_('abc').select('foo').force_index('egg')
+
+        self.assertEqual('SELECT "foo" FROM "abc" FORCE INDEX ("egg")', str(q))
+
+    def test_select_with_force_index_with_index_object(self):
+        index = Index('egg')
+        q = Query.from_('abc').select('foo').force_index(index)
+
+        self.assertEqual('SELECT "foo" FROM "abc" FORCE INDEX ("egg")', str(q))
+
+    def test_select_with_force_index_multiple_indexes(self):
+        q = Query.from_('abc').select('foo').force_index('egg', 'bacon')
+
+        self.assertEqual('SELECT "foo" FROM "abc" FORCE INDEX ("egg","bacon")', str(q))
+
+    def test_select_with_force_index_multiple_calls(self):
+        q = Query.from_('abc').select('foo').force_index('egg',).force_index('spam')
+
+        self.assertEqual('SELECT "foo" FROM "abc" FORCE INDEX ("egg","spam")', str(q))
+
     def test_mysql_query_uses_backtick_quote_chars(self):
         q = MySQLQuery.from_('abc').select('foo', 'bar')
 
@@ -343,6 +365,10 @@ class WhereTests(unittest.TestCase):
 
         self.assertEqual('SELECT * FROM "abc"', str(q1))
 
+    def test_select_with_force_index_and_where(self):
+        q = Query.from_('abc').select('foo').where(self.t.foo == self.t.bar).force_index('egg')
+
+        self.assertEqual('SELECT "foo" FROM "abc" FORCE INDEX ("egg") WHERE "foo"="bar"', str(q))
 
 class PreWhereTests(WhereTests):
     t = Table('abc')