Repository: kayak/pypika
Pull request #300 (2019-08-20 15:32:16)

Support simple SQL Server TOP clauses
This PR implements supports for simple SQL Server `TOP` clauses, which limit the number of rows returned by a query.  Official documentation on the syntax is here: https://docs.microsoft.com/en-us/sql/t-sql/queries/top-transact-sql?view=sql-server-2017.

The basic idea is that after your `SELECT` keyword, you specify `TOP (<number>)`, where `<number>` specifies that max number of rows to retrieve; for example, `SELECT TOP(1) xyz FROM abc ORDER BY xyz DESC`. The parentheses are optional but recommended, so I've included them.

The `PERCENT` and `WITH TIES` arguments aren't supported.  I can take a look at those if that would be helpful/important.

Thanks!


Patch:
diff --git a/pypika/dialects.py b/pypika/dialects.py
index 17fa9a6..5336f34 100644
--- a/pypika/dialects.py
+++ b/pypika/dialects.py
@@ -303,10 +303,39 @@ class RedshiftQuery(Query):
 class MSSQLQueryBuilder(QueryBuilder):
     def __init__(self):
         super(MSSQLQueryBuilder, self).__init__(dialect=Dialects.MSSQL)
+        self._top = None
+
+    @builder
+    def top(self, value):
+        """
+        Implements support for simple TOP clauses.
+
+        Does not include support for PERCENT or WITH TIES.
+
+        https://docs.microsoft.com/en-us/sql/t-sql/queries/top-transact-sql?view=sql-server-2017
+        """
+        try:
+            self._top = int(value)
+        except ValueError:
+            raise QueryException('TOP value must be an integer')
 
     def get_sql(self, *args, **kwargs):
         return super(MSSQLQueryBuilder, self).get_sql(*args, groupby_alias=False, **kwargs)
 
+    def _top_sql(self):
+        if self._top:
+            return 'TOP ({}) '.format(self._top)
+        else:
+            return ''
+
+    def _select_sql(self, **kwargs):
+        return 'SELECT {distinct}{top}{select}'.format(
+            top=self._top_sql(),
+            distinct='DISTINCT ' if self._distinct else '',
+            select=','.join(term.get_sql(with_alias=True, subquery=True, **kwargs)
+                            for term in self._selects),
+        )
+
 
 class MSSQLQuery(Query):
     """


Test patch:
diff --git a/pypika/tests/dialects/test_mssql.py b/pypika/tests/dialects/test_mssql.py
new file mode 100644
index 0000000..5ded243
--- /dev/null
+++ b/pypika/tests/dialects/test_mssql.py
@@ -0,0 +1,33 @@
+import unittest
+
+from pypika import Table
+from pypika.dialects import MSSQLQuery
+from pypika.utils import QueryException
+
+
+class SelectTests(unittest.TestCase):
+
+    def test_normal_select(self):
+        q = MSSQLQuery.from_('abc').select('def')
+
+        self.assertEqual('SELECT "def" FROM "abc"', str(q))
+
+    def test_distinct_select(self):
+        q = MSSQLQuery.from_('abc').select('def').distinct()
+
+        self.assertEqual('SELECT DISTINCT "def" FROM "abc"', str(q))
+
+    def test_top_distinct_select(self):
+        q = MSSQLQuery.from_('abc').select('def').top(10).distinct()
+
+        self.assertEqual('SELECT DISTINCT TOP (10) "def" FROM "abc"', str(q))
+
+    def test_top_select(self):
+        q = MSSQLQuery.from_('abc').select('def').top(10)
+
+        self.assertEqual('SELECT TOP (10) "def" FROM "abc"', str(q))
+
+    def test_top_select_non_int(self):
+        with self.assertRaisesRegex(QueryException,
+                                    'TOP value must be an integer'):
+            MSSQLQuery.from_('abc').select('def').top('a')