From f3fcd1fdebb21b0bffeb2e9b2c101e33abf504df Mon Sep 17 00:00:00 2001 From: Dave Bauman Date: Fri, 14 Jun 2024 10:41:28 -0400 Subject: [PATCH] feat: Apply row limit transform to query in backend --- .../server/datasources/query_execution.py | 11 ++++++- querybook/server/logic/admin.py | 12 +++++++ querybook/server/tasks/run_datadoc.py | 11 +++++++ .../DataDocQueryCell/DataDocQueryCell.tsx | 12 +++---- .../QueryComposer/QueryComposer.tsx | 8 +++-- .../components/QueryComposer/RunQuery.tsx | 32 +++++++------------ querybook/webapp/redux/dataDoc/action.ts | 3 ++ .../webapp/redux/queryExecutions/action.ts | 2 ++ querybook/webapp/resource/queryExecution.ts | 8 ++++- 9 files changed, 68 insertions(+), 31 deletions(-) diff --git a/querybook/server/datasources/query_execution.py b/querybook/server/datasources/query_execution.py index 8c3515cdb..42b30d792 100644 --- a/querybook/server/datasources/query_execution.py +++ b/querybook/server/datasources/query_execution.py @@ -13,6 +13,7 @@ ) 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, @@ -61,10 +62,18 @@ @register("/query_execution/", methods=["POST"]) -def create_query_execution(query, engine_id, data_cell_id=None, originator=None): +def create_query_execution( + query, engine_id, row_limit=-1, 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 diff --git a/querybook/server/logic/admin.py b/querybook/server/logic/admin.py index c570041db..b08be57d4 100644 --- a/querybook/server/logic/admin.py +++ b/querybook/server/logic/admin.py @@ -187,6 +187,18 @@ def get_admin_announcements(session=None): ) +@with_session +def get_engine_feature_param( + engine_id, feature_param_name, default_value=None, session=None +): + query_engine = get_query_engine_by_id(engine_id, session=session) + return ( + query_engine.get_feature_params().get(feature_param_name, default_value) + if query_engine + else default_value + ) + + """ --------------------------------------------------------------------------------------------------------- QUERY METASTORE ? diff --git a/querybook/server/tasks/run_datadoc.py b/querybook/server/tasks/run_datadoc.py index 0f109759a..abf0821cb 100644 --- a/querybook/server/tasks/run_datadoc.py +++ b/querybook/server/tasks/run_datadoc.py @@ -11,10 +11,12 @@ from lib.logger import get_logger from lib.query_analysis.templating import render_templated_query +from lib.query_analysis.transform import transform_to_limited_query from lib.scheduled_datadoc.export import export_datadoc from lib.scheduled_datadoc.legacy import convert_if_legacy_datadoc_schedule from lib.scheduled_datadoc.notification import notifiy_on_datadoc_complete +from logic import admin as admin_logic from logic import datadoc as datadoc_logic from logic import query_execution as qe_logic from logic.schedule import ( @@ -73,6 +75,7 @@ def run_datadoc_with_config( # Prepping chain jobs each unit is a [make_qe_task, run_query_task] combo for index, query_cell in enumerate(query_cells): engine_id = query_cell.meta["engine"] + limit = query_cell.meta.get("limit", -1) raw_query = query_cell.context # Skip empty cells @@ -86,6 +89,14 @@ def run_datadoc_with_config( engine_id, session=session, ) + + # If meta["limit"] is set and > 0, apply limit to the query + row_limit_enabled = admin_logic.get_engine_feature_param( + engine_id, "row_limit", False, session=session + ) + if row_limit_enabled and limit >= 0: + query = transform_to_limited_query(query, limit) + except Exception as e: on_datadoc_completion( is_success=False, diff --git a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx index acc873680..2afbe17c0 100644 --- a/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx +++ b/querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx @@ -45,12 +45,7 @@ import { import { DEFAULT_ROW_LIMIT } from 'lib/sql-helper/sql-limiter'; import { getPossibleTranspilers } from 'lib/templated-query/transpile'; import { enableResizable } from 'lib/utils'; -import { - getShortcutSymbols, - KeyMap, - matchKeyMap, - matchKeyPress, -} from 'lib/utils/keyboard'; +import { getShortcutSymbols, KeyMap, matchKeyPress } from 'lib/utils/keyboard'; import { doesLanguageSupportUDF } from 'lib/utils/udf'; import * as dataDocActions from 'redux/dataDoc/action'; import * as dataSourcesActions from 'redux/dataSources/action'; @@ -479,6 +474,7 @@ class DataDocQueryCellComponent extends React.PureComponent { await this.props.createQueryExecution( query, engineId, + this.rowLimit, this.props.cellId ) ).id; @@ -550,6 +546,7 @@ class DataDocQueryCellComponent extends React.PureComponent { return this.props.createQueryExecution( renderedQuery, this.engineId, + this.rowLimit, this.props.cellId ); } @@ -1073,8 +1070,9 @@ function mapDispatchToProps(dispatch: Dispatch) { createQueryExecution: ( query: string, engineId: number, + rowLimit: number, cellId: number - ) => dispatch(createQueryExecution(query, engineId, cellId)), + ) => dispatch(createQueryExecution(query, engineId, rowLimit, cellId)), setTableSidebarId: (id: number) => dispatch(setSidebarTableId(id)), diff --git a/querybook/webapp/components/QueryComposer/QueryComposer.tsx b/querybook/webapp/components/QueryComposer/QueryComposer.tsx index a3e213396..5d1686ffc 100644 --- a/querybook/webapp/components/QueryComposer/QueryComposer.tsx +++ b/querybook/webapp/components/QueryComposer/QueryComposer.tsx @@ -41,7 +41,7 @@ import { useTrackView } from 'hooks/useTrackView'; import { trackClick } from 'lib/analytics'; import { createSQLLinter } from 'lib/codemirror/codemirror-lint'; import { replaceStringIndices, searchText } from 'lib/data-doc/search'; -import { getSelectedQuery, IRange, TableToken } from 'lib/sql-helper/sql-lexer'; +import { getSelectedQuery, IRange } from 'lib/sql-helper/sql-lexer'; import { DEFAULT_ROW_LIMIT } from 'lib/sql-helper/sql-limiter'; import { getPossibleTranspilers } from 'lib/templated-query/transpile'; import { enableResizable, getQueryEngineId, sleep } from 'lib/utils'; @@ -519,7 +519,11 @@ const QueryComposer: React.FC = () => { engine.id, async (query, engineId) => { const data = await dispatch( - queryExecutionsAction.createQueryExecution(query, engineId) + queryExecutionsAction.createQueryExecution( + query, + engineId, + rowLimit + ) ); return data.id; } diff --git a/querybook/webapp/components/QueryComposer/RunQuery.tsx b/querybook/webapp/components/QueryComposer/RunQuery.tsx index 5b67daec5..8cf801b82 100644 --- a/querybook/webapp/components/QueryComposer/RunQuery.tsx +++ b/querybook/webapp/components/QueryComposer/RunQuery.tsx @@ -5,10 +5,7 @@ import { ISamplingTables, TDataDocMetaVariables } from 'const/datadoc'; import { IQueryEngine } from 'const/queryEngine'; import { sendConfirm } from 'lib/querybookUI'; import { getDroppedTables } from 'lib/sql-helper/sql-checker'; -import { - getLimitedQuery, - hasQueryContainUnlimitedSelect, -} from 'lib/sql-helper/sql-limiter'; +import { hasQueryContainUnlimitedSelect } from 'lib/sql-helper/sql-limiter'; import { renderTemplatedQuery } from 'lib/templated-query'; import { Nullable } from 'lib/typescript'; import { formatError } from 'lib/utils/error'; @@ -46,13 +43,9 @@ export async function transformQuery( sampleRate ); - const limitedQuery = await transformLimitedQuery( - sampledQuery, - rowLimit, - engine - ); + await checkUnlimitedQuery(sampledQuery, rowLimit, engine); - return limitedQuery; + return sampledQuery; } export async function runQuery( @@ -87,17 +80,16 @@ async function transformTemplatedQuery( } } -async function transformLimitedQuery( +async function checkUnlimitedQuery( query: string, rowLimit: Nullable, engine: IQueryEngine ) { - if (!engine.feature_params?.row_limit) { - return query; - } - - if (rowLimit != null && rowLimit >= 0) { - return getLimitedQuery(query, rowLimit, engine.language); + if ( + !engine.feature_params?.row_limit || + (rowLimit != null && rowLimit >= 0) + ) { + return; } // query is unlimited but engine has row limit feature turned on @@ -108,11 +100,11 @@ async function transformLimitedQuery( ); if (!unlimitedSelectQuery) { - return query; + return; } // Show a warning modal to let user confirm what they are doing - return new Promise((resolve, reject) => { + return new Promise((resolve, reject) => { sendConfirm({ header: 'Your SELECT query is unbounded', message: ( @@ -135,7 +127,7 @@ async function transformLimitedQuery( ), - onConfirm: () => resolve(query), + onConfirm: () => resolve(), onDismiss: () => reject(), confirmText: 'Run without LIMIT', }); diff --git a/querybook/webapp/redux/dataDoc/action.ts b/querybook/webapp/redux/dataDoc/action.ts index 7886bb1c9..ce08c7a39 100644 --- a/querybook/webapp/redux/dataDoc/action.ts +++ b/querybook/webapp/redux/dataDoc/action.ts @@ -26,6 +26,7 @@ import { } from 'lib/data-doc/datadoc-permission'; import dataDocSocket from 'lib/data-doc/datadoc-socketio'; import { convertRawToContentState } from 'lib/richtext/serialize'; +import { DEFAULT_ROW_LIMIT } from 'lib/sql-helper/sql-limiter'; import { getQueryEngineId } from 'lib/utils'; import { DataDocAccessRequestResource, @@ -308,9 +309,11 @@ export function insertDataDocCell( queryEngineIds ); const engine = meta?.['engine'] ?? defaultQueryEngine; + const limit = meta?.['limit'] ?? DEFAULT_ROW_LIMIT; meta = { ...meta, engine, + limit, }; } diff --git a/querybook/webapp/redux/queryExecutions/action.ts b/querybook/webapp/redux/queryExecutions/action.ts index d1e365504..1fef726f9 100644 --- a/querybook/webapp/redux/queryExecutions/action.ts +++ b/querybook/webapp/redux/queryExecutions/action.ts @@ -378,6 +378,7 @@ export function pollQueryExecution( export function createQueryExecution( query: string, engineId?: number, + rowLimit?: number, cellId?: number ): ThunkResult> { return async (dispatch, getState) => { @@ -387,6 +388,7 @@ export function createQueryExecution( const { data: queryExecution } = await QueryExecutionResource.create( query, selectedEngineId, + rowLimit, cellId ); dispatch(receiveQueryExecution(queryExecution, cellId)); diff --git a/querybook/webapp/resource/queryExecution.ts b/querybook/webapp/resource/queryExecution.ts index e595f077b..27b1adf0d 100644 --- a/querybook/webapp/resource/queryExecution.ts +++ b/querybook/webapp/resource/queryExecution.ts @@ -66,10 +66,16 @@ export const QueryExecutionResource = { environment_id: environmentId, }), - create: (query: string, engineId: number, cellId?: number) => { + create: ( + query: string, + engineId: number, + rowLimit?: number, + cellId?: number + ) => { const params = { query, engine_id: engineId, + row_limit: rowLimit, }; if (cellId != null) {