You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@superset.apache.org by GitBox <gi...@apache.org> on 2018/05/31 18:12:49 UTC

[GitHub] timifasubaa closed pull request #5023: [sqllab] force limit queries only when there is no existing limit

timifasubaa closed pull request #5023: [sqllab] force limit queries only when there is no existing limit
URL: https://github.com/apache/incubator-superset/pull/5023
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py
index 0e189b1259..cd5ab87226 100644
--- a/superset/db_engine_specs.py
+++ b/superset/db_engine_specs.py
@@ -104,15 +104,32 @@ def apply_limit_to_sql(cls, sql, limit, database):
             )
             return database.compile_sqla_query(qry)
         elif LimitMethod.FORCE_LIMIT:
-            no_limit = re.sub(r"""
+            sql_without_limit = cls.get_query_without_limit(sql)
+            return '{sql_without_limit} LIMIT {limit}'.format(**locals())
+        return sql
+
+    @classmethod
+    def get_limit_from_sql(cls, sql):
+        limit_pattern = re.compile(r"""
+                (?ix)          # case insensitive, verbose
+                \s+            # whitespace
+                LIMIT\s+(\d+)  # LIMIT $ROWS
+                ;?             # optional semi-colon
+                (\s|;)*$       # remove trailing spaces tabs or semicolons
+                """)
+        matches = limit_pattern.findall(sql)
+        if matches:
+            return int(matches[0][0])
+
+    @classmethod
+    def get_query_without_limit(cls, sql):
+        return re.sub(r"""
                 (?ix)        # case insensitive, verbose
                 \s+          # whitespace
                 LIMIT\s+\d+  # LIMIT $ROWS
                 ;?           # optional semi-colon
                 (\s|;)*$     # remove trailing spaces tabs or semicolons
                 """, '', sql)
-            return '{no_limit} LIMIT {limit}'.format(**locals())
-        return sql
 
     @staticmethod
     def csv_to_df(**kwargs):
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 68e3f70fee..c9f07ae906 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -171,6 +171,7 @@ def handle_error(msg):
     # Limit enforced only for retrieving the data, not for the CTA queries.
     superset_query = SupersetQuery(rendered_query)
     executed_sql = superset_query.stripped()
+    SQL_MAX_ROWS = app.config.get('SQL_MAX_ROW')
     if not superset_query.is_select() and not database.allow_dml:
         return handle_error(
             'Only `SELECT` statements are allowed against this database')
@@ -185,9 +186,10 @@ def handle_error(msg):
                 query.user_id, start_dttm.strftime('%Y_%m_%d_%H_%M_%S'))
         executed_sql = superset_query.as_create_table(query.tmp_table_name)
         query.select_as_cta_used = True
-    elif (query.limit and superset_query.is_select()):
+    if (superset_query.is_select() and SQL_MAX_ROWS and
+            (not query.limit or query.limit > SQL_MAX_ROWS)):
+        query.limit = SQL_MAX_ROWS
         executed_sql = database.apply_limit_to_sql(executed_sql, query.limit)
-        query.limit_used = True
 
     # Hook to allow environment-specific mutation (usually comments) to the SQL
     SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
diff --git a/superset/views/core.py b/superset/views/core.py
index 40d24b268a..b4a1689a91 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -2394,7 +2394,7 @@ def sql_json(self):
 
         query = Query(
             database_id=int(database_id),
-            limit=int(app.config.get('SQL_MAX_ROW', None)),
+            limit=mydb.db_engine_spec.get_limit_from_sql(sql),
             sql=sql,
             schema=schema,
             select_as_cta=request.form.get('select_as_cta') == 'true',
diff --git a/tests/celery_tests.py b/tests/celery_tests.py
index f6d1a2958f..39b7749ae8 100644
--- a/tests/celery_tests.py
+++ b/tests/celery_tests.py
@@ -196,10 +196,9 @@ def test_run_async_query(self):
         self.assertEqual([{'name': 'Admin'}], df.to_dict(orient='records'))
         self.assertEqual(QueryStatus.SUCCESS, query.status)
         self.assertTrue('FROM tmp_async_1' in query.select_sql)
-        self.assertTrue('LIMIT 666' in query.select_sql)
         self.assertEqual(
             'CREATE TABLE tmp_async_1 AS \nSELECT name FROM ab_role '
-            "WHERE name='Admin'", query.executed_sql)
+            "WHERE name='Admin' LIMIT 666", query.executed_sql)
         self.assertEqual(sql_where, query.sql)
         self.assertEqual(0, query.rows)
         self.assertEqual(666, query.limit)
@@ -207,6 +206,33 @@ def test_run_async_query(self):
         self.assertEqual(True, query.select_as_cta)
         self.assertEqual(True, query.select_as_cta_used)
 
+    def test_run_async_query_with_lower_limit(self):
+        main_db = self.get_main_database(db.session)
+        eng = main_db.get_sqla_engine()
+        sql_where = "SELECT name FROM ab_role WHERE name='Alpha' LIMIT 1"
+        result = self.run_sql(
+            main_db.id, sql_where, '5', async='true', tmp_table='tmp_async_2',
+            cta='true')
+        assert result['query']['state'] in (
+            QueryStatus.PENDING, QueryStatus.RUNNING, QueryStatus.SUCCESS)
+
+        time.sleep(1)
+
+        query = self.get_query_by_id(result['query']['serverId'])
+        df = pd.read_sql_query(query.select_sql, con=eng)
+        self.assertEqual(QueryStatus.SUCCESS, query.status)
+        self.assertEqual([{'name': 'Alpha'}], df.to_dict(orient='records'))
+        self.assertEqual(QueryStatus.SUCCESS, query.status)
+        self.assertTrue('FROM tmp_async_2' in query.select_sql)
+        self.assertEqual(
+            'CREATE TABLE tmp_async_2 AS \nSELECT name FROM ab_role '
+            "WHERE name='Alpha' LIMIT 1", query.executed_sql)
+        self.assertEqual(sql_where, query.sql)
+        self.assertEqual(0, query.rows)
+        self.assertEqual(1, query.limit)
+        self.assertEqual(True, query.select_as_cta)
+        self.assertEqual(True, query.select_as_cta_used)
+
     @staticmethod
     def de_unicode_dict(d):
         def str_if_basestring(o):
diff --git a/tests/db_engine_specs_test.py b/tests/db_engine_specs_test.py
index c38e4f5690..bdce0b060d 100644
--- a/tests/db_engine_specs_test.py
+++ b/tests/db_engine_specs_test.py
@@ -95,6 +95,19 @@ def sql_limit_regex(
         limited = engine_spec_class.apply_limit_to_sql(sql, limit, main)
         self.assertEquals(expected_sql, limited)
 
+    def test_extract_limit_from_query(self, engine_spec_class=MySQLEngineSpec):
+        q0 = 'select * from table'
+        q1 = 'select * from mytable limit 10'
+        q2 = 'select * from (select * from my_subquery limit 10) where col=1 limit 20'
+        q3 = 'select * from (select * from my_subquery limit 10);'
+        q4 = 'select * from (select * from my_subquery limit 10) where col=1 limit 20;'
+
+        self.assertEqual(engine_spec_class.get_limit_from_sql(q0), None)
+        self.assertEqual(engine_spec_class.get_limit_from_sql(q1), 10)
+        self.assertEqual(engine_spec_class.get_limit_from_sql(q2), 20)
+        self.assertEqual(engine_spec_class.get_limit_from_sql(q3), None)
+        self.assertEqual(engine_spec_class.get_limit_from_sql(q4), 20)
+
     def test_wrapped_query(self):
         self.sql_limit_regex(
             'SELECT * FROM a',


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org