From d8be0a7dd50dce8c27079a8d1a1c36376c8f353b Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Mon, 13 May 2019 00:34:34 -0500
Subject: [PATCH] Break line before LIMIT statement to prevent trailing comment
 issue (#7485)

* Break line before LIMIT statement to prevent trailing comment issue

This may not be a perfect solution but it addresses the issue in 7483

closes https://github.com/apache/incubator-superset/issues/7483

* fix tests
---
 superset/sql_parse.py         |  2 +-
 tests/celery_tests.py         |  3 ++-
 tests/db_engine_specs_test.py |  4 ++--
 tests/sql_parse_tests.py      | 19 +++++++++++++++++++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index 662f6c326..63ae05e66 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -168,7 +168,7 @@ class ParsedQuery(object):
         """returns the query with the specified limit"""
         """does not change the underlying query"""
         if not self._limit:
-            return self.sql + ' LIMIT ' + str(new_limit)
+            return f'{self.sql}\nLIMIT {new_limit}'
         limit_pos = None
         tokens = self._parsed[0].tokens
         # Add all items to before_str until there is a limit
diff --git a/tests/celery_tests.py b/tests/celery_tests.py
index 80b410186..bccffb211 100644
--- a/tests/celery_tests.py
+++ b/tests/celery_tests.py
@@ -201,7 +201,8 @@ class CeleryTestCase(SupersetTestCase):
         self.assertEqual(
             'CREATE TABLE tmp_async_1 AS \n'
             'SELECT name FROM ab_role '
-            "WHERE name='Admin' LIMIT 666", query.executed_sql)
+            "WHERE name='Admin'\n"
+            'LIMIT 666', query.executed_sql)
         self.assertEqual(sql_where, query.sql)
         self.assertEqual(0, query.rows)
         self.assertEqual(False, query.limit_used)
diff --git a/tests/db_engine_specs_test.py b/tests/db_engine_specs_test.py
index 2a22c590e..e0d914f0b 100644
--- a/tests/db_engine_specs_test.py
+++ b/tests/db_engine_specs_test.py
@@ -190,7 +190,7 @@ class DbEngineSpecsTestCase(SupersetTestCase):
     def test_simple_limit_query(self):
         self.sql_limit_regex(
             'SELECT * FROM a',
-            'SELECT * FROM a LIMIT 1000',
+            'SELECT * FROM a\nLIMIT 1000',
         )
 
     def test_modify_limit_query(self):
@@ -288,7 +288,7 @@ class DbEngineSpecsTestCase(SupersetTestCase):
                     'LIMIT 777'""",
             """
                 SELECT
-                    'LIMIT 777' LIMIT 1000""",
+                    'LIMIT 777'\nLIMIT 1000""",
         )
 
     def test_time_grain_blacklist(self):
diff --git a/tests/sql_parse_tests.py b/tests/sql_parse_tests.py
index 7096147b5..fb3dc7894 100644
--- a/tests/sql_parse_tests.py
+++ b/tests/sql_parse_tests.py
@@ -431,6 +431,25 @@ class SupersetTestCase(unittest.TestCase):
         """
         self.assertEquals({'SalesOrderHeader'}, self.extract_tables(query))
 
+    def test_get_query_with_new_limit_comment(self):
+        sql = 'SELECT * FROM ab_user --SOME COMMENT'
+        parsed = sql_parse.ParsedQuery(sql)
+        newsql = parsed.get_query_with_new_limit(1000)
+        self.assertEquals(newsql, sql + '\nLIMIT 1000')
+
+    def test_get_query_with_new_limit_comment_with_limit(self):
+        sql = 'SELECT * FROM ab_user --SOME COMMENT WITH LIMIT 555'
+        parsed = sql_parse.ParsedQuery(sql)
+        newsql = parsed.get_query_with_new_limit(1000)
+        self.assertEquals(newsql, sql + '\nLIMIT 1000')
+
+    def test_get_query_with_new_limit(self):
+        sql = 'SELECT * FROM ab_user LIMIT 555'
+        parsed = sql_parse.ParsedQuery(sql)
+        newsql = parsed.get_query_with_new_limit(1000)
+        expected = 'SELECT * FROM ab_user LIMIT 1000'
+        self.assertEquals(newsql, expected)
+
     def test_basic_breakdown_statements(self):
         multi_sql = """
         SELECT * FROM ab_user;