From 32eef059b621f87fe9e037a702ac216611638657 Mon Sep 17 00:00:00 2001
From: Elkin <elkin@raketa.im>
Date: Wed, 29 Jan 2020 14:16:53 +0300
Subject: [PATCH] MSSQL 2012 OFFSET/FETCH support

---
 lib/sqlalchemy/dialects/mssql/base.py | 44 +++++++++++++++++++++++++--
 test/dialect/mssql/test_compiler.py   | 23 ++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index 94f6a3303a..cb601218c7 100644
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -1635,8 +1635,42 @@ class MSSQLCompiler(compiler.SQLCompiler):
         return text
 
     def limit_clause(self, select, **kw):
-        # Limit in mssql is after the select keyword
-        return ""
+        """ MSSQL 2012 supports OFFSET/FETCH operators
+            Use it instead subquery with row_number
+
+        """
+
+        if self.dialect._supports_offset_fetch and (
+                (
+                        not select._simple_int_limit
+                        and select._limit_clause is not None
+                )
+                or (
+                        select._offset_clause is not None
+                        and not select._simple_int_offset
+                        or select._offset
+                )
+        ):
+            # OFFSET are FETCH are options of the ORDER BY clause
+            if not select._order_by_clause.clauses:
+                raise exc.CompileError(
+                    "MSSQL requires an order_by when "
+                    "using an OFFSET or a non-simple "
+                    "LIMIT clause"
+                )
+
+            text = ""
+            if select._offset_clause is not None:
+                text += "\n OFFSET %s ROWS" % self.process(
+                    select._offset_clause, **kw
+                )
+            if select._limit_clause is not None:
+                text += "\n FETCH NEXT %s ROWS ONLY " % self.process(
+                    select._limit_clause, **kw
+                )
+            return text
+        else:
+            return ""
 
     def visit_try_cast(self, element, **kw):
         return "TRY_CAST (%s AS %s)" % (
@@ -1647,9 +1681,10 @@ class MSSQLCompiler(compiler.SQLCompiler):
     def visit_select(self, select, **kwargs):
         """Look for ``LIMIT`` and OFFSET in a select statement, and if
         so tries to wrap it in a subquery with ``row_number()`` criterion.
+        MSSQL 2012 and above are excluded
 
         """
-        if (
+        if not self.dialect._supports_offset_fetch and (
             (not select._simple_int_limit and select._limit_clause is not None)
             or (
                 select._offset_clause is not None
@@ -2288,6 +2323,7 @@ class MSDialect(default.DefaultDialect):
     non_native_boolean_check_constraint = False
     supports_unicode_binds = True
     postfetch_lastrowid = True
+    _supports_offset_fetch = False
 
     server_version_info = ()
 
@@ -2436,6 +2472,8 @@ class MSDialect(default.DefaultDialect):
                 self.server_version_info >= MS_2012_VERSION
             )
 
+        self._supports_offset_fetch = (self.server_version_info and self.server_version_info[0] >= 12)
+
     def _get_default_schema_name(self, connection):
         if self.server_version_info < MS_2005_VERSION:
             return self.schema_name
diff --git a/test/dialect/mssql/test_compiler.py b/test/dialect/mssql/test_compiler.py
index bb5199b00d..07ceb5bf5b 100644
--- a/test/dialect/mssql/test_compiler.py
+++ b/test/dialect/mssql/test_compiler.py
@@ -844,6 +844,29 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL):
         assert t.c.x in set(c._create_result_map()["x"][1])
         assert t.c.y in set(c._create_result_map()["y"][1])
 
+    def test_limit_offset_using_offset_fetch(self):
+        t = table("t", column("x", Integer), column("y", Integer))
+        dialect_2012 = mssql.base.MSDialect()
+        dialect_2012._supports_offset_fetch = True
+
+        s = select([t]).where(t.c.x == 5).order_by(t.c.y).limit(10).offset(20)
+
+        self.assert_compile(
+            s,
+            "SELECT t.x, t.y "
+            "FROM t "
+            "WHERE t.x = :x_1 ORDER BY t.y "
+            "OFFSET :param_1 ROWS "
+            "FETCH NEXT :param_2 ROWS ONLY ",
+            checkparams={"param_1": 20, "param_2": 10, "x_1": 5},
+            dialect=dialect_2012
+        )
+
+        c = s.compile(dialect=dialect_2012)
+        eq_(len(c._result_columns), 2)
+        assert t.c.x in set(c._create_result_map()["x"][1])
+        assert t.c.y in set(c._create_result_map()["y"][1])
+
     def test_limit_offset_w_ambiguous_cols(self):
         t = table("t", column("x", Integer), column("y", Integer))
 
-- 
2.47.3