Repository: kayak/pypika
Pull request #478 (2020-09-02 09:24:10)

Support FETCH NEXT N ROWS ONLY statement in MSSQLQueryBuilder
* I re-used some limit/offset code for MSSQL offset/fetch, as this is the de facto way now of limiting/offsetting rows in modern versions of SQL Server. There are more configuration options, but we can look at supporting them if/when needed. What is provided here is for common use cases.
* I also provided a `fetch_next` method which is has the same behaviour as using the `limit` method to make it more friendly to T-SQL users.


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index bf77f5e..a356818 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -578,9 +578,31 @@ class MSSQLQueryBuilder(QueryBuilder):
         except ValueError:
             raise QueryException("TOP value must be an integer")
 
+    @builder
+    def fetch_next(self, limit: int) -> "MSSQLQueryBuilder":
+        # Overridden to provide a more domain-specific API for T-SQL users
+        self._limit = limit
+
+    def _offset_sql(self) -> str:
+        return " OFFSET {offset} ROWS".format(offset=self._offset or 0)
+
+    def _limit_sql(self) -> str:
+        return " FETCH NEXT {limit} ROWS ONLY".format(limit=self._limit)
+
+    def _apply_pagination(self, querystring: str) -> str:
+        # Note: Overridden as MSSQL specifies offset before the fetch next limit
+        if self._limit is not None or self._offset:
+            # Offset has to be present if fetch next is specified in a MSSQL query
+            querystring += self._offset_sql()
+
+        if self._limit is not None:
+            querystring += self._limit_sql()
+
+        return querystring
+
     def get_sql(self, *args: Any, **kwargs: Any) -> str:
         return super().get_sql(
-              *args, groupby_alias=False, **kwargs
+            *args, groupby_alias=False, **kwargs
         )
 
     def _top_sql(self) -> str:
diff --git a/pypika/queries.py b/pypika/queries.py
index 0810e40..d14aaae 100644
--- a/pypika/queries.py
+++ b/pypika/queries.py
@@ -594,12 +594,10 @@ class QueryBuilder(Selectable, Term):
         self._with = []
         self._selects = []
         self._force_indexes = []
-        self._use_indexes = []
         self._columns = []
         self._values = []
         self._distinct = False
         self._ignore = False
-        self._for_update = False
 
         self._wheres = None
         self._prewheres = None
@@ -830,22 +828,10 @@ 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
 
-    @builder
-    def for_update(self) -> "QueryBuilder":
-        self._for_update = True
-
     @builder
     def ignore(self) -> "QueryBuilder":
         self._ignore = True
@@ -1251,9 +1237,6 @@ 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
@@ -1276,14 +1259,7 @@ class QueryBuilder(Selectable, Term):
         if self._orderbys:
             querystring += self._orderby_sql(**kwargs)
 
-        if self._limit is not None:
-            querystring += self._limit_sql()
-
-        if self._offset:
-            querystring += self._offset_sql()
-
-        if self._for_update:
-            querystring += self._for_update_sql()
+        querystring = self._apply_pagination(querystring)
 
         if subquery:
             querystring = "({query})".format(query=querystring)
@@ -1296,13 +1272,22 @@ class QueryBuilder(Selectable, Term):
 
         return querystring
 
+    def _apply_pagination(self, querystring: str) -> str:
+        if self._limit is not None:
+            querystring += self._limit_sql()
+
+        if self._offset:
+            querystring += self._offset_sql()
+
+        return querystring
+
     def _with_sql(self, **kwargs: Any) -> str:
         return "WITH " + ",".join(
-              clause.name
-              + " AS ("
-              + clause.get_sql(subquery=False, with_alias=False, **kwargs)
-              + ") "
-              for clause in self._with
+            clause.name
+            + " AS ("
+            + clause.get_sql(subquery=False, with_alias=False, **kwargs)
+            + ") "
+            for clause in self._with
         )
 
     def _distinct_sql(self, **kwargs: Any) -> str:
@@ -1313,14 +1298,6 @@ class QueryBuilder(Selectable, Term):
 
         return distinct
 
-    def _for_update_sql(self) -> str:
-        if self._for_update:
-            for_update = ' FOR UPDATE'
-        else:
-            for_update = ''
-
-        return for_update
-
     def _select_sql(self, **kwargs: Any) -> str:
         return "SELECT {distinct}{select}".format(
               distinct=self._distinct_sql(**kwargs),
@@ -1389,11 +1366,6 @@ 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/dialects/test_mssql.py b/pypika/tests/dialects/test_mssql.py
index 625ed97..2c8dd4a 100644
--- a/pypika/tests/dialects/test_mssql.py
+++ b/pypika/tests/dialects/test_mssql.py
@@ -28,3 +28,23 @@ class SelectTests(unittest.TestCase):
     def test_top_select_non_int(self):
         with self.assertRaisesRegex(QueryException, "TOP value must be an integer"):
             MSSQLQuery.from_("abc").select("def").top("a")
+
+    def test_limit(self):
+        q = MSSQLQuery.from_("abc").select("def").orderby("def").limit(10)
+
+        self.assertEqual('SELECT "def" FROM "abc" ORDER BY "def" OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY', str(q))
+
+    def test_fetch_next(self):
+        q = MSSQLQuery.from_("abc").select("def").orderby("def").fetch_next(10)
+
+        self.assertEqual('SELECT "def" FROM "abc" ORDER BY "def" OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY', str(q))
+
+    def test_offset(self):
+        q = MSSQLQuery.from_("abc").select("def").orderby("def").offset(10)
+
+        self.assertEqual('SELECT "def" FROM "abc" ORDER BY "def" OFFSET 10 ROWS', str(q))
+
+    def test_fetch_next_with_offset(self):
+        q = MSSQLQuery.from_("abc").select("def").orderby("def").fetch_next(10).offset(10)
+
+        self.assertEqual('SELECT "def" FROM "abc" ORDER BY "def" OFFSET 10 ROWS FETCH NEXT 10 ROWS ONLY', str(q))
diff --git a/pypika/tests/test_selects.py b/pypika/tests/test_selects.py
index a140b43..c629585 100644
--- a/pypika/tests/test_selects.py
+++ b/pypika/tests/test_selects.py
@@ -236,27 +236,6 @@ 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")
 
@@ -315,10 +294,6 @@ class WhereTests(unittest.TestCase):
         q = self.t.select("*").where(self.t.foo == self.t.bar)
         self.assertEqual(q, q1)
 
-    def test_where_field_equals_for_update(self):
-        q = Query.from_(self.t).select("*").where(self.t.foo == self.t.bar).for_update()
-        self.assertEqual('SELECT * FROM "abc" WHERE "foo"="bar" FOR UPDATE', str(q))
-
     def test_where_field_equals_where(self):
         q = (
             Query.from_(self.t)