fix: PR comments

querybook/server/datasources/query_execution.py

@@ -13,7 +13,6 @@
 )
 from clients.common import FileDoesNotExist
 from lib.export.all_exporters import ALL_EXPORTERS, get_exporter
-from lib.query_analysis.transform import transform_to_limited_query
 from lib.result_store import GenericReader
 from lib.query_analysis.templating import (
     QueryTemplatingError,
@@ -62,18 +61,10 @@
 
 
 @register("/query_execution/", methods=["POST"])
-def create_query_execution(
-    query, engine_id, row_limit=-1, data_cell_id=None, originator=None
-):
+def create_query_execution(query, engine_id, data_cell_id=None, originator=None):
     with DBSession() as session:
         verify_query_engine_permission(engine_id, session=session)
 
-        row_limit_enabled = admin_logic.get_engine_feature_param(
-            engine_id, "row_limit", False, session=session
-        )
-        if row_limit_enabled and row_limit >= 0:
-            query = transform_to_limited_query(query, row_limit)
-
         uid = current_user.id
         query_execution = logic.create_query_execution(
             query=query, engine_id=engine_id, uid=uid, session=session

querybook/server/datasources/query_transform.py

@@ -1,8 +1,30 @@
 from app.datasource import register
+from lib.logger import get_logger
 from lib.query_analysis.transform import (
+    has_query_contains_unlimited_select,
+    transform_to_limited_query,
     transform_to_sampled_query,
 )
 
+LOG = get_logger(__file__)
+
+
+@register("/query/transform/limited/", methods=["POST"])
+def query_limited(
+    query: str,
+    row_limit: int,
+    language: str,
+):
+    limited_query = transform_to_limited_query(
+        query=query, limit=row_limit, language=language
+    )
+
+    unlimited_select = has_query_contains_unlimited_select(
+        query=limited_query, language=language
+    )
+
+    return {"query": limited_query, "unlimited_select": unlimited_select}
+
 
 @register("/query/transform/sampling/", methods=["POST"])
 def query_sampling(

querybook/server/lib/query_analysis/transform.py

@@ -67,24 +67,34 @@ def get_limited_select_statement(statement_ast: exp.Expression, limit: int):
     return statement_ast.limit(limit)
 
 
-def has_query_contains_unlimited_select(query: str, language: str) -> bool:
+def has_query_contains_unlimited_select(query: str, language: str = None):
     """Check if a query contains a select statement without a limit.
     Args:
         query: The query to check
     Returns:
-        bool: True if the query contains a select statement without a limit, False otherwise
+        str: The first select statement without a limit. None if all select statements have a limit.
     """
-    statements = parse(query, dialect=_get_sqlglot_dialect[language])
-    return any(get_select_statement_limit(s) == -1 for s in statements)
+    dialect = _get_sqlglot_dialect(language)
+    statements = parse(query, dialect)
+    return next(
+        (
+            s.sql(dialect=dialect, pretty=True)
+            for s in statements
+            if get_select_statement_limit(s) == -1
+        ),
+        None,
+    )
 
 
 def transform_to_limited_query(
     query: str, limit: int = None, language: str = None
 ) -> str:
     """Apply a limit to all select statements in a query if they don't already have a limit.
     It returns a new query with the limit applied and the original query is not modified.
+
+    If limit is None or negative, the query is returned as-is.
     """
-    if not limit:
+    if not limit or limit < 0:
         return query
 
     try:

querybook/tests/test_lib/test_query_analysis/test_transform.py

@@ -3,6 +3,7 @@
 from lib.query_analysis.transform import (
     format_query,
     get_select_statement_limit,
+    has_query_contains_unlimited_select,
     transform_to_limited_query,
     transform_to_sampled_query,
 )
@@ -46,6 +47,50 @@ def test_select_with_limit(self):
                 self.assertEqual(get_select_statement_limit(query), expected)
 
 
+class HasQueryContainsUnlimitedSelectTestCase(TestCase):
+    def test_select_limit(self):
+        tests = [
+            "SELECT 1 LIMIT 10",
+            "SELECT * FROM table_1 WHERE field = 1 LIMIT 10",
+            "TRUNCATE TABLE table_1; SELECT * FROM table_1 WHERE field = 1 LIMIT 1000",
+            "SELECT * FROM table_1 WHERE field = 1 LIMIT 10; SELECT * FROM table_2 WHERE field = 1 LIMIT 1000",
+        ]
+        for query in tests:
+            with self.subTest(query=query):
+                self.assertIsNone(has_query_contains_unlimited_select(query))
+
+    def test_select_no_limit(self):
+        tests = [
+            ("SELECT 1", "SELECT\n  1"),
+            (
+                "SELECT * FROM table_1 WHERE field = 1",
+                "SELECT\n  *\nFROM table_1\nWHERE\n  field = 1",
+            ),
+            ("SELECT 1; SELECT 2", "SELECT\n  1"),
+            (
+                "SELECT * FROM table_1 WHERE field = 1 LIMIT 10; SELECT * FROM table_1 WHERE field = 1",
+                "SELECT\n  *\nFROM table_1\nWHERE\n  field = 1",
+            ),
+        ]
+        for query, expected in tests:
+            with self.subTest(query=query):
+                self.assertEquals(has_query_contains_unlimited_select(query), expected)
+
+    def test_not_select_statements(self):
+        tests = [
+            "DELETE FROM table_1 WHERE field = 1",
+            "CREATE DATABASE IF NOT EXISTS db_1",
+            "CREATE TABLE table_1 (field1 INT)",
+            "TRUNCATE TABLE table_1",
+            "DROP TABLE IF EXISTS db.table1; CREATE TABLE db.table1",
+            "INSERT INTO table_1 (field1) VALUES (1)",
+            "UPDATE table_1 SET field1 = 1 WHERE field = 1",
+        ]
+        for query in tests:
+            with self.subTest(query=query):
+                self.assertIsNone(has_query_contains_unlimited_select(query))
+
+
 class GetLimitedQueryTestCase(TestCase):
     def test_limit_is_not_specified(self):
         tests = [

querybook/webapp/__tests__/lib/sql-helper/sql-limiter.test.ts

@@ -1,7 +1,4 @@
-import {
-    getLimitedQuery,
-    getSelectStatementLimit,
-} from 'lib/sql-helper/sql-limiter';
+import { getSelectStatementLimit } from 'lib/sql-helper/sql-limiter';
 
 describe('getSelectStatementLimit', () => {
     describe('when it is not a SELECT statement', () => {
@@ -76,126 +73,3 @@ describe('getSelectStatementLimit', () => {
         });
     });
 });
-
-describe('getLimitedQuery', () => {
-    describe('not limited', () => {
-        test('when rowLimit is not specified', () => {
-            const query = 'SELECT * FROM table_1 WHERE field = 1;';
-            expect(getLimitedQuery(query)).toBe(query);
-        });
-        test('when rowLimit is not specified for multiple queries', () => {
-            const query = `
-                SELECT * FROM table_1 WHERE field = 1;
-                SELECT * FROM table_1 WHERE field = 1;
-            `;
-            expect(getLimitedQuery(query)).toBe(query);
-        });
-        test('when running a select query with fetch', () => {
-            const query =
-                'SELECT * FROM table_1 ORDER BY id FETCH FIRST 10 ROWS ONLY;';
-            expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
-        });
-        test('when running a select query with offset and fetch', () => {
-            const query =
-                'SELECT * FROM table_1 ORDER BY id OFFSET 10 FETCH NEXT 10 ROWS ONLY;';
-            expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
-        });
-        test('when running a select query with nested query', () => {
-            const query = `select * from (select * from table limit 5) as x limit 10`;
-            expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
-        });
-        test('when running a select query with a where clause and a limit', () => {
-            const query = 'SELECT * FROM table_1 WHERE field = 1 LIMIT 1000;';
-            expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
-        });
-    });
-    describe('limited', () => {
-        test('when running a select query', () => {
-            const query = 'SELECT * FROM table_1';
-            expect(getLimitedQuery(query, 10)).toBe(`${query} limit 10;`);
-        });
-        test('when running a select query with a where clause and a group by and an order by', () => {
-            const query =
-                'SELECT field, count(*) FROM table_1 WHERE deleted = false GROUP BY field ORDER BY field';
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with trailing semicolon', () => {
-            const query = 'SELECT * FROM table_1;';
-            expect(getLimitedQuery(query, 10)).toBe(
-                'SELECT * FROM table_1 limit 10;'
-            );
-        });
-        test('when running a select query with comments', () => {
-            const query = 'SELECT * FROM table_1 -- limit here';
-            expect(getLimitedQuery(query, 10)).toBe(
-                'SELECT * FROM table_1 limit 10;'
-            );
-        });
-        test('when running a select query with non-keyword limits', () => {
-            const query = `SELECT id, account, 'limit' FROM querybook2.limit ORDER by 'limit' ASC`;
-            expect(getLimitedQuery(query, 10)).toBe(`${query} limit 10;`);
-        });
-        test('when running a multiple select queries', () => {
-            const query = `SELECT * FROM table_1;
-SELECT col1, col2, FROM table2;`;
-            expect(getLimitedQuery(query, 10)).toBe(
-                `SELECT * FROM table_1 limit 10;
-SELECT col1, col2, FROM table2 limit 10;`
-            );
-        });
-        test('when running a select query with a where clause', () => {
-            const query = 'SELECT * FROM table_1 WHERE field = 1';
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with a where clause and an order by', () => {
-            const query =
-                'SELECT * FROM table_1 WHERE field = 1 ORDER BY field';
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with a where clause and a group by and an order by', () => {
-            const query =
-                'SELECT field, count(*) FROM table_1 WHERE deleted = false GROUP BY field ORDER BY field';
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running two select queries with mixed limits', () => {
-            const query = `SELECT * FROM table_1;
-SELECT col1, col2, FROM table2 LIMIT 300;`;
-            expect(getLimitedQuery(query, 10))
-                .toBe(`SELECT * FROM table_1 limit 10;
-SELECT col1, col2, FROM table2 LIMIT 300;`);
-        });
-        test('when running multiple select queries with mixed limits', () => {
-            const query = `SELECT * FROM table_1;
-SELECT col1, col2, FROM table2 LIMIT 300;
-SELECT field, count(1) FROM table3 GROUP BY field`;
-            expect(getLimitedQuery(query, 10))
-                .toBe(`SELECT * FROM table_1 limit 10;
-SELECT col1, col2, FROM table2 LIMIT 300;
-SELECT field, count(1) FROM table3 GROUP BY field limit 10;`);
-        });
-        test('when running a select query with nested query', () => {
-            const query = `select * from (select * from table limit 5) as x`;
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with wrapped where', () => {
-            const query = `select * from table where (field = 1 and field2 = 2)`;
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with two nested queries', () => {
-            const query = `select * from (select * from table limit 5) as x outer join (select * from table2 limit 5) as y on x.id = y.id`;
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with two nested queries', () => {
-            const query = `select * from (select * from table limit 5) as x outer join (select * from table2 limit 5) as y on x.id = y.id`;
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with two union queries', () => {
-            const query = `select id, name from table_a union all select id, name from table_b where (deleted = false and active = true)`;
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-        test('when running a select query with two nested union queries', () => {
-            const query = `(select id, name from table_a limit 10) union all (select id, name from table_b where (deleted = false and active = true))`;
-            expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
-        });
-    });
-});

querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx

@@ -477,7 +477,6 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                     await this.props.createQueryExecution(
                         query,
                         engineId,
-                        this.rowLimit,
                         this.props.cellId
                     )
                 ).id;
@@ -549,7 +548,6 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
             return this.props.createQueryExecution(
                 renderedQuery,
                 this.engineId,
-                this.rowLimit,
                 this.props.cellId
             );
         }
@@ -1094,9 +1092,8 @@ function mapDispatchToProps(dispatch: Dispatch) {
         createQueryExecution: (
             query: string,
             engineId: number,
-            rowLimit: number,
             cellId: number
-        ) => dispatch(createQueryExecution(query, engineId, rowLimit, cellId)),
+        ) => dispatch(createQueryExecution(query, engineId, cellId)),
 
         setTableSidebarId: (id: number) => dispatch(setSidebarTableId(id)),
 

querybook/webapp/components/QueryComposer/QueryComposer.tsx

@@ -524,11 +524,7 @@ const QueryComposer: React.FC = () => {
             engine.id,
             async (query, engineId) => {
                 const data = await dispatch(
-                    queryExecutionsAction.createQueryExecution(
-                        query,
-                        engineId,
-                        rowLimit
-                    )
+                    queryExecutionsAction.createQueryExecution(query, engineId)
                 );
                 return data.id;
             }