Skip to content

Commit b8a33e0

Browse files
baumandmbaumandm
committed
fix: PR comments
1 parent eff3c2f commit b8a33e0

File tree

13 files changed

+118
-232
lines changed

13 files changed

+118
-232
lines changed

‎querybook/server/datasources/query_execution.py

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
)
1414
from clients.common import FileDoesNotExist
1515
from lib.export.all_exporters import ALL_EXPORTERS, get_exporter
16-
from lib.query_analysis.transform import transform_to_limited_query
1716
from lib.result_store import GenericReader
1817
from lib.query_analysis.templating import (
1918
QueryTemplatingError,
@@ -62,18 +61,10 @@
6261

6362

6463
@register("/query_execution/", methods=["POST"])
65-
def create_query_execution(
66-
query, engine_id, row_limit=-1, data_cell_id=None, originator=None
67-
):
64+
def create_query_execution(query, engine_id, data_cell_id=None, originator=None):
6865
with DBSession() as session:
6966
verify_query_engine_permission(engine_id, session=session)
7067

71-
row_limit_enabled = admin_logic.get_engine_feature_param(
72-
engine_id, "row_limit", False, session=session
73-
)
74-
if row_limit_enabled and row_limit >= 0:
75-
query = transform_to_limited_query(query, row_limit)
76-
7768
uid = current_user.id
7869
query_execution = logic.create_query_execution(
7970
query=query, engine_id=engine_id, uid=uid, session=session

‎querybook/server/datasources/query_transform.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,30 @@
11
from app.datasource import register
2+
from lib.logger import get_logger
23
from lib.query_analysis.transform import (
4+
has_query_contains_unlimited_select,
5+
transform_to_limited_query,
36
transform_to_sampled_query,
47
)
58

9+
LOG = get_logger(__file__)
10+
11+
12+
@register("/query/transform/limited/", methods=["POST"])
13+
def query_limited(
14+
query: str,
15+
row_limit: int,
16+
language: str,
17+
):
18+
limited_query = transform_to_limited_query(
19+
query=query, limit=row_limit, language=language
20+
)
21+
22+
unlimited_select = has_query_contains_unlimited_select(
23+
query=limited_query, language=language
24+
)
25+
26+
return {"query": limited_query, "unlimited_select": unlimited_select}
27+
628

729
@register("/query/transform/sampling/", methods=["POST"])
830
def query_sampling(

‎querybook/server/lib/query_analysis/transform.py

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -67,24 +67,34 @@ def get_limited_select_statement(statement_ast: exp.Expression, limit: int):
6767
return statement_ast.limit(limit)
6868

6969

70-
def has_query_contains_unlimited_select(query: str, language: str) -> bool:
70+
def has_query_contains_unlimited_select(query: str, language: str = None):
7171
"""Check if a query contains a select statement without a limit.
7272
Args:
7373
query: The query to check
7474
Returns:
75-
bool: True if the query contains a select statement without a limit, False otherwise
75+
str: The first select statement without a limit. None if all select statements have a limit.
7676
"""
77-
statements = parse(query, dialect=_get_sqlglot_dialect[language])
78-
return any(get_select_statement_limit(s) == -1 for s in statements)
77+
dialect = _get_sqlglot_dialect(language)
78+
statements = parse(query, dialect)
79+
return next(
80+
(
81+
s.sql(dialect=dialect, pretty=True)
82+
for s in statements
83+
if get_select_statement_limit(s) == -1
84+
),
85+
None,
86+
)
7987

8088

8189
def transform_to_limited_query(
8290
query: str, limit: int = None, language: str = None
8391
) -> str:
8492
"""Apply a limit to all select statements in a query if they don't already have a limit.
8593
It returns a new query with the limit applied and the original query is not modified.
94+
95+
If limit is None or negative, the query is returned as-is.
8696
"""
87-
if not limit:
97+
if not limit or limit < 0:
8898
return query
8999

90100
try:

‎querybook/tests/test_lib/test_query_analysis/test_transform.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from lib.query_analysis.transform import (
44
format_query,
55
get_select_statement_limit,
6+
has_query_contains_unlimited_select,
67
transform_to_limited_query,
78
transform_to_sampled_query,
89
)
@@ -46,6 +47,50 @@ def test_select_with_limit(self):
4647
self.assertEqual(get_select_statement_limit(query), expected)
4748

4849

50+
class HasQueryContainsUnlimitedSelectTestCase(TestCase):
51+
def test_select_limit(self):
52+
tests = [
53+
"SELECT 1 LIMIT 10",
54+
"SELECT * FROM table_1 WHERE field = 1 LIMIT 10",
55+
"TRUNCATE TABLE table_1; SELECT * FROM table_1 WHERE field = 1 LIMIT 1000",
56+
"SELECT * FROM table_1 WHERE field = 1 LIMIT 10; SELECT * FROM table_2 WHERE field = 1 LIMIT 1000",
57+
]
58+
for query in tests:
59+
with self.subTest(query=query):
60+
self.assertIsNone(has_query_contains_unlimited_select(query))
61+
62+
def test_select_no_limit(self):
63+
tests = [
64+
("SELECT 1", "SELECT\n 1"),
65+
(
66+
"SELECT * FROM table_1 WHERE field = 1",
67+
"SELECT\n *\nFROM table_1\nWHERE\n field = 1",
68+
),
69+
("SELECT 1; SELECT 2", "SELECT\n 1"),
70+
(
71+
"SELECT * FROM table_1 WHERE field = 1 LIMIT 10; SELECT * FROM table_1 WHERE field = 1",
72+
"SELECT\n *\nFROM table_1\nWHERE\n field = 1",
73+
),
74+
]
75+
for query, expected in tests:
76+
with self.subTest(query=query):
77+
self.assertEquals(has_query_contains_unlimited_select(query), expected)
78+
79+
def test_not_select_statements(self):
80+
tests = [
81+
"DELETE FROM table_1 WHERE field = 1",
82+
"CREATE DATABASE IF NOT EXISTS db_1",
83+
"CREATE TABLE table_1 (field1 INT)",
84+
"TRUNCATE TABLE table_1",
85+
"DROP TABLE IF EXISTS db.table1; CREATE TABLE db.table1",
86+
"INSERT INTO table_1 (field1) VALUES (1)",
87+
"UPDATE table_1 SET field1 = 1 WHERE field = 1",
88+
]
89+
for query in tests:
90+
with self.subTest(query=query):
91+
self.assertIsNone(has_query_contains_unlimited_select(query))
92+
93+
4994
class GetLimitedQueryTestCase(TestCase):
5095
def test_limit_is_not_specified(self):
5196
tests = [

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

Lines changed: 1 addition & 127 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,4 @@
1-
import {
2-
getLimitedQuery,
3-
getSelectStatementLimit,
4-
} from 'lib/sql-helper/sql-limiter';
1+
import { getSelectStatementLimit } from 'lib/sql-helper/sql-limiter';
52

63
describe('getSelectStatementLimit', () => {
74
describe('when it is not a SELECT statement', () => {
@@ -76,126 +73,3 @@ describe('getSelectStatementLimit', () => {
7673
});
7774
});
7875
});
79-
80-
describe('getLimitedQuery', () => {
81-
describe('not limited', () => {
82-
test('when rowLimit is not specified', () => {
83-
const query = 'SELECT * FROM table_1 WHERE field = 1;';
84-
expect(getLimitedQuery(query)).toBe(query);
85-
});
86-
test('when rowLimit is not specified for multiple queries', () => {
87-
const query = `
88-
SELECT * FROM table_1 WHERE field = 1;
89-
SELECT * FROM table_1 WHERE field = 1;
90-
`;
91-
expect(getLimitedQuery(query)).toBe(query);
92-
});
93-
test('when running a select query with fetch', () => {
94-
const query =
95-
'SELECT * FROM table_1 ORDER BY id FETCH FIRST 10 ROWS ONLY;';
96-
expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
97-
});
98-
test('when running a select query with offset and fetch', () => {
99-
const query =
100-
'SELECT * FROM table_1 ORDER BY id OFFSET 10 FETCH NEXT 10 ROWS ONLY;';
101-
expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
102-
});
103-
test('when running a select query with nested query', () => {
104-
const query = `select * from (select * from table limit 5) as x limit 10`;
105-
expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
106-
});
107-
test('when running a select query with a where clause and a limit', () => {
108-
const query = 'SELECT * FROM table_1 WHERE field = 1 LIMIT 1000;';
109-
expect(getLimitedQuery(query, 100, 'trino')).toBe(query);
110-
});
111-
});
112-
describe('limited', () => {
113-
test('when running a select query', () => {
114-
const query = 'SELECT * FROM table_1';
115-
expect(getLimitedQuery(query, 10)).toBe(`${query} limit 10;`);
116-
});
117-
test('when running a select query with a where clause and a group by and an order by', () => {
118-
const query =
119-
'SELECT field, count(*) FROM table_1 WHERE deleted = false GROUP BY field ORDER BY field';
120-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
121-
});
122-
test('when running a select query with trailing semicolon', () => {
123-
const query = 'SELECT * FROM table_1;';
124-
expect(getLimitedQuery(query, 10)).toBe(
125-
'SELECT * FROM table_1 limit 10;'
126-
);
127-
});
128-
test('when running a select query with comments', () => {
129-
const query = 'SELECT * FROM table_1 -- limit here';
130-
expect(getLimitedQuery(query, 10)).toBe(
131-
'SELECT * FROM table_1 limit 10;'
132-
);
133-
});
134-
test('when running a select query with non-keyword limits', () => {
135-
const query = `SELECT id, account, 'limit' FROM querybook2.limit ORDER by 'limit' ASC`;
136-
expect(getLimitedQuery(query, 10)).toBe(`${query} limit 10;`);
137-
});
138-
test('when running a multiple select queries', () => {
139-
const query = `SELECT * FROM table_1;
140-
SELECT col1, col2, FROM table2;`;
141-
expect(getLimitedQuery(query, 10)).toBe(
142-
`SELECT * FROM table_1 limit 10;
143-
SELECT col1, col2, FROM table2 limit 10;`
144-
);
145-
});
146-
test('when running a select query with a where clause', () => {
147-
const query = 'SELECT * FROM table_1 WHERE field = 1';
148-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
149-
});
150-
test('when running a select query with a where clause and an order by', () => {
151-
const query =
152-
'SELECT * FROM table_1 WHERE field = 1 ORDER BY field';
153-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
154-
});
155-
test('when running a select query with a where clause and a group by and an order by', () => {
156-
const query =
157-
'SELECT field, count(*) FROM table_1 WHERE deleted = false GROUP BY field ORDER BY field';
158-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
159-
});
160-
test('when running two select queries with mixed limits', () => {
161-
const query = `SELECT * FROM table_1;
162-
SELECT col1, col2, FROM table2 LIMIT 300;`;
163-
expect(getLimitedQuery(query, 10))
164-
.toBe(`SELECT * FROM table_1 limit 10;
165-
SELECT col1, col2, FROM table2 LIMIT 300;`);
166-
});
167-
test('when running multiple select queries with mixed limits', () => {
168-
const query = `SELECT * FROM table_1;
169-
SELECT col1, col2, FROM table2 LIMIT 300;
170-
SELECT field, count(1) FROM table3 GROUP BY field`;
171-
expect(getLimitedQuery(query, 10))
172-
.toBe(`SELECT * FROM table_1 limit 10;
173-
SELECT col1, col2, FROM table2 LIMIT 300;
174-
SELECT field, count(1) FROM table3 GROUP BY field limit 10;`);
175-
});
176-
test('when running a select query with nested query', () => {
177-
const query = `select * from (select * from table limit 5) as x`;
178-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
179-
});
180-
test('when running a select query with wrapped where', () => {
181-
const query = `select * from table where (field = 1 and field2 = 2)`;
182-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
183-
});
184-
test('when running a select query with two nested queries', () => {
185-
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`;
186-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
187-
});
188-
test('when running a select query with two nested queries', () => {
189-
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`;
190-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
191-
});
192-
test('when running a select query with two union queries', () => {
193-
const query = `select id, name from table_a union all select id, name from table_b where (deleted = false and active = true)`;
194-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
195-
});
196-
test('when running a select query with two nested union queries', () => {
197-
const query = `(select id, name from table_a limit 10) union all (select id, name from table_b where (deleted = false and active = true))`;
198-
expect(getLimitedQuery(query, 100)).toBe(`${query} limit 100;`);
199-
});
200-
});
201-
});

‎querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -477,7 +477,6 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
477477
await this.props.createQueryExecution(
478478
query,
479479
engineId,
480-
this.rowLimit,
481480
this.props.cellId
482481
)
483482
).id;
@@ -549,7 +548,6 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
549548
return this.props.createQueryExecution(
550549
renderedQuery,
551550
this.engineId,
552-
this.rowLimit,
553551
this.props.cellId
554552
);
555553
}
@@ -1094,9 +1092,8 @@ function mapDispatchToProps(dispatch: Dispatch) {
10941092
createQueryExecution: (
10951093
query: string,
10961094
engineId: number,
1097-
rowLimit: number,
10981095
cellId: number
1099-
) => dispatch(createQueryExecution(query, engineId, rowLimit, cellId)),
1096+
) => dispatch(createQueryExecution(query, engineId, cellId)),
11001097

11011098
setTableSidebarId: (id: number) => dispatch(setSidebarTableId(id)),
11021099

‎querybook/webapp/components/QueryComposer/QueryComposer.tsx

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -524,11 +524,7 @@ const QueryComposer: React.FC = () => {
524524
engine.id,
525525
async (query, engineId) => {
526526
const data = await dispatch(
527-
queryExecutionsAction.createQueryExecution(
528-
query,
529-
engineId,
530-
rowLimit
531-
)
527+
queryExecutionsAction.createQueryExecution(query, engineId)
532528
);
533529
return data.id;
534530
}

0 commit comments

Comments
 (0)