feat: Apply row limit transform to query in backend

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

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 ?

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,

querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx

@@ -477,6 +477,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                     await this.props.createQueryExecution(
                         query,
                         engineId,
+                        this.rowLimit,
                         this.props.cellId
                     )
                 ).id;
@@ -548,6 +549,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
             return this.props.createQueryExecution(
                 renderedQuery,
                 this.engineId,
+                this.rowLimit,
                 this.props.cellId
             );
         }
@@ -1092,8 +1094,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)),
 

querybook/webapp/components/QueryComposer/QueryComposer.tsx

@@ -42,7 +42,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';
@@ -524,7 +524,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;
             }

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<number>,
     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<string>((resolve, reject) => {
+    return new Promise<void>((resolve, reject) => {
         sendConfirm({
             header: 'Your SELECT query is unbounded',
             message: (
@@ -135,7 +127,7 @@ async function transformLimitedQuery(
                     </pre>
                 </Content>
             ),
-            onConfirm: () => resolve(query),
+            onConfirm: () => resolve(),
             onDismiss: () => reject(),
             confirmText: 'Run without LIMIT',
         });

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,
             };
         }
 

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<Promise<IQueryExecution>> {
     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));

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) {