diff --git a/superset/sql_parse.py b/superset/sql_parse.py index 662f6c32622..63ae05e6616 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 80b41018683..bccffb21144 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 2a22c590ed6..e0d914f0b51 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 7096147b561..fb3dc7894bd 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;