Add limit option for MSSQL query runner (#6704)

* Add limit option for MSSQL query runner

* Fixed linting errors

redash/query_runner/__init__.py

@@ -119,6 +119,7 @@ class BaseQueryRunner:
     noop_query = None
     limit_query = " LIMIT 1000"
     limit_keywords = ["LIMIT", "OFFSET"]
+    limit_after_select = False
 
     def __init__(self, configuration):
         self.syntax = "sql"
@@ -301,10 +302,19 @@ def add_limit_to_query(self, query):
         parsed_query = sqlparse.parse(query)[0]
         limit_tokens = sqlparse.parse(self.limit_query)[0].tokens
         length = len(parsed_query.tokens)
-        if parsed_query.tokens[length - 1].ttype == sqlparse.tokens.Punctuation:
-            parsed_query.tokens[length - 1 : length - 1] = limit_tokens
+        if not self.limit_after_select:
+            if parsed_query.tokens[length - 1].ttype == sqlparse.tokens.Punctuation:
+                parsed_query.tokens[length - 1 : length - 1] = limit_tokens
+            else:
+                parsed_query.tokens += limit_tokens
         else:
-            parsed_query.tokens += limit_tokens
+            for i in range(length - 1, -1, -1):
+                if parsed_query[i].value.upper() == "SELECT":
+                    index = parsed_query.token_index(parsed_query[i + 1])
+                    parsed_query = sqlparse.sql.Statement(
+                        parsed_query.tokens[:index] + limit_tokens + parsed_query.tokens[index:]
+                    )
+                    break
         return str(parsed_query)
 
     def apply_auto_limit(self, query_text, should_apply_auto_limit):

redash/query_runner/mssql.py

@@ -34,6 +34,10 @@ class SqlServer(BaseSQLQueryRunner):
     should_annotate_query = False
     noop_query = "SELECT 1"
 
+    limit_query = " TOP 1000"
+    limit_keywords = ["TOP"]
+    limit_after_select = True
+
     @classmethod
     def configuration_schema(cls):
         return {

redash/query_runner/mssql_odbc.py

@@ -21,6 +21,10 @@ class SQLServerODBC(BaseSQLQueryRunner):
     should_annotate_query = False
     noop_query = "SELECT 1"
 
+    limit_query = " TOP 1000"
+    limit_keywords = ["TOP"]
+    limit_after_select = True
+
     @classmethod
     def configuration_schema(cls):
         return {