From e853a978288c3479e7e78f947c79f4dc7db0cab6 Mon Sep 17 00:00:00 2001 From: Vadim Ogievetsky Date: Mon, 29 Jan 2024 20:09:46 -0800 Subject: [PATCH] Web console: Make table driven query modification actions work with slices. (#15779) * Make table driven query modification actions work with slices. * cleanup found query prefix * fix regex complexity --- web-console/src/utils/sql.spec.ts | 70 +++++++++---- web-console/src/utils/sql.ts | 17 +++- .../workbench-view/query-tab/query-tab.tsx | 97 ++++++++++++------- .../result-table-pane/result-table-pane.tsx | 41 ++++---- 4 files changed, 154 insertions(+), 71 deletions(-) diff --git a/web-console/src/utils/sql.spec.ts b/web-console/src/utils/sql.spec.ts index 9f107b1c45fc..cdb6fabd8004 100644 --- a/web-console/src/utils/sql.spec.ts +++ b/web-console/src/utils/sql.spec.ts @@ -85,6 +85,7 @@ describe('sql', () => { "column": 14, "row": 1, }, + "index": 0, "sql": "SELECT * FROM wikipedia", "startOffset": 0, @@ -99,6 +100,7 @@ describe('sql', () => { "column": 7, "row": 5, }, + "index": 1, "sql": "SELECT * FROM w2 LIMIT 5", @@ -132,6 +134,7 @@ describe('sql', () => { "column": 15, "row": 5, }, + "index": 0, "sql": "SELECT \\"channel\\", COUNT(*) AS \\"Count\\" @@ -150,6 +153,7 @@ describe('sql', () => { "column": 31, "row": 3, }, + "index": 1, "sql": "SELECT * FROM \\"wikipedia\\"", "startOffset": 48, "startRowColumn": Object { @@ -184,6 +188,7 @@ describe('sql', () => { "column": 15, "row": 8, }, + "index": 0, "sql": "WITH w1 AS ( SELECT channel, page FROM \\"wikipedia\\" ) @@ -205,6 +210,7 @@ describe('sql', () => { "column": 39, "row": 1, }, + "index": 1, "sql": "SELECT channel, page FROM \\"wikipedia\\"", "startOffset": 15, "startRowColumn": Object { @@ -218,6 +224,7 @@ describe('sql', () => { "column": 15, "row": 8, }, + "index": 2, "sql": "SELECT page, COUNT(*) AS \\"cnt\\" @@ -234,8 +241,10 @@ describe('sql', () => { `); }); - it('works with replace query', () => { + it('works with select query followed by a replace query', () => { const text = sane` + SELECT * FROM "wiki" + REPLACE INTO "wikipedia" OVERWRITE ALL WITH "ext" AS ( SELECT * @@ -259,11 +268,26 @@ describe('sql', () => { expect(found).toMatchInlineSnapshot(` Array [ Object { - "endOffset": 379, + "endOffset": 29, + "endRowColumn": Object { + "column": 7, + "row": 2, + }, + "index": 0, + "sql": "SELECT * FROM \\"wiki\\"", + "startOffset": 0, + "startRowColumn": Object { + "column": 0, + "row": 0, + }, + }, + Object { + "endOffset": 401, "endRowColumn": Object { "column": 18, - "row": 15, + "row": 17, }, + "index": 1, "sql": "REPLACE INTO \\"wikipedia\\" OVERWRITE ALL WITH \\"ext\\" AS ( SELECT * @@ -280,18 +304,19 @@ describe('sql', () => { \\"channel\\" FROM \\"ext\\" PARTITIONED BY DAY", - "startOffset": 0, + "startOffset": 22, "startRowColumn": Object { "column": 0, - "row": 0, + "row": 2, }, }, Object { - "endOffset": 360, + "endOffset": 382, "endRowColumn": Object { "column": 10, - "row": 14, + "row": 16, }, + "index": 2, "sql": "WITH \\"ext\\" AS ( SELECT * FROM TABLE( @@ -306,18 +331,19 @@ describe('sql', () => { \\"isRobot\\", \\"channel\\" FROM \\"ext\\"", - "startOffset": 39, + "startOffset": 61, "startRowColumn": Object { "column": 0, - "row": 1, + "row": 3, }, }, Object { - "endOffset": 276, + "endOffset": 298, "endRowColumn": Object { "column": 70, - "row": 8, + "row": 10, }, + "index": 3, "sql": "SELECT * FROM TABLE( EXTERN( @@ -325,27 +351,28 @@ describe('sql', () => { '{\\"type\\":\\"json\\"}' ) ) EXTEND (\\"isRobot\\" VARCHAR, \\"channel\\" VARCHAR, \\"timestamp\\" VARCHAR)", - "startOffset": 57, + "startOffset": 79, "startRowColumn": Object { "column": 2, - "row": 2, + "row": 4, }, }, Object { - "endOffset": 360, + "endOffset": 382, "endRowColumn": Object { "column": 10, - "row": 14, + "row": 16, }, + "index": 4, "sql": "SELECT TIME_PARSE(\\"timestamp\\") AS \\"__time\\", \\"isRobot\\", \\"channel\\" FROM \\"ext\\"", - "startOffset": 279, + "startOffset": 301, "startRowColumn": Object { "column": 0, - "row": 10, + "row": 12, }, }, ] @@ -384,6 +411,7 @@ describe('sql', () => { "column": 22, "row": 17, }, + "index": 0, "sql": "EXPLAIN PLAN FOR INSERT INTO \\"wikipedia\\" WITH \\"ext\\" AS ( @@ -414,6 +442,7 @@ describe('sql', () => { "column": 22, "row": 17, }, + "index": 1, "sql": "INSERT INTO \\"wikipedia\\" WITH \\"ext\\" AS ( SELECT * @@ -443,6 +472,7 @@ describe('sql', () => { "column": 10, "row": 15, }, + "index": 2, "sql": "WITH \\"ext\\" AS ( SELECT * FROM TABLE( @@ -469,6 +499,7 @@ describe('sql', () => { "column": 70, "row": 9, }, + "index": 3, "sql": "SELECT * FROM TABLE( EXTERN( @@ -488,6 +519,7 @@ describe('sql', () => { "column": 10, "row": 15, }, + "index": 4, "sql": "SELECT TIME_PARSE(\\"timestamp\\") AS \\"__time\\", \\"isRobot\\", @@ -526,6 +558,7 @@ describe('sql', () => { "column": 14, "row": 2, }, + "index": 0, "sql": "EXPLAIN PLAN FOR SELECT * FROM wikipedia", @@ -541,6 +574,7 @@ describe('sql', () => { "column": 14, "row": 2, }, + "index": 1, "sql": "SELECT * FROM wikipedia", "startOffset": 17, @@ -555,6 +589,7 @@ describe('sql', () => { "column": 7, "row": 7, }, + "index": 2, "sql": "EXPLAIN PLAN FOR SELECT * FROM w2 @@ -571,6 +606,7 @@ describe('sql', () => { "column": 7, "row": 7, }, + "index": 3, "sql": "SELECT * FROM w2 LIMIT 5", diff --git a/web-console/src/utils/sql.ts b/web-console/src/utils/sql.ts index b80c5ea457d5..78e69a75a18e 100644 --- a/web-console/src/utils/sql.ts +++ b/web-console/src/utils/sql.ts @@ -108,7 +108,21 @@ export function findSqlQueryPrefix(text: string): string | undefined { } } +export function cleanSqlQueryPrefix(text: string): string { + const matchReplace = text.match(/\sREPLACE$/i); + if (matchReplace) { + // This query likely grabbed a "REPLACE" (which is not a reserved keyword) from the next query over, see if we can delete it + const textWithoutReplace = text.slice(0, -matchReplace[0].length).trimEnd(); + if (SqlQuery.maybeParse(textWithoutReplace)) { + return textWithoutReplace; + } + } + + return text; +} + export interface QuerySlice { + index: number; startOffset: number; startRowColumn: RowColumn; endOffset: number; @@ -130,11 +144,12 @@ export function findAllSqlQueriesInText(text: string): QuerySlice[] { if (sql) { const endIndex = m.index + sql.length; found.push({ + index: found.length, startOffset: offset + m.index, startRowColumn: offsetToRowColumn(text, offset + m.index)!, endOffset: offset + endIndex, endRowColumn: offsetToRowColumn(text, offset + endIndex)!, - sql, + sql: cleanSqlQueryPrefix(sql), }); } remainingText = remainingText.slice(advanceBy); diff --git a/web-console/src/views/workbench-view/query-tab/query-tab.tsx b/web-console/src/views/workbench-view/query-tab/query-tab.tsx index 8a4129fc67bf..8c7db0427365 100644 --- a/web-console/src/views/workbench-view/query-tab/query-tab.tsx +++ b/web-console/src/views/workbench-view/query-tab/query-tab.tsx @@ -44,6 +44,7 @@ import { WorkbenchRunningPromises } from '../../../singletons/workbench-running- import type { ColumnMetadata, QueryAction, QuerySlice, RowColumn } from '../../../utils'; import { DruidError, + findAllSqlQueriesInText, localStorageGet, LocalStorageKeys, localStorageSet, @@ -122,15 +123,40 @@ export const QueryTab = React.memo(function QueryTab(props: QueryTabProps) { onQueryChange(query.changeQueryString(queryString)); }); - const parsedQuery = query.getParsedQuery(); - const handleQueryAction = usePermanentCallback((queryAction: QueryAction) => { - if (!(parsedQuery instanceof SqlQuery)) return; - onQueryChange(query.changeQueryString(parsedQuery.apply(queryAction).toString())); - - if (shouldAutoRun()) { - setTimeout(() => void handleRun(false), 20); - } - }); + const handleQueryAction = usePermanentCallback( + (queryAction: QueryAction, sliceIndex: number | undefined) => { + let newQueryString: string; + if (typeof sliceIndex === 'number') { + const { queryString } = query; + const foundQuery = findAllSqlQueriesInText(queryString)[sliceIndex]; + if (!foundQuery) return; + const parsedQuery = SqlQuery.maybeParse(foundQuery.sql); + if (!parsedQuery) return; + newQueryString = + queryString.slice(0, foundQuery.startOffset) + + parsedQuery.apply(queryAction) + + queryString.slice(foundQuery.endOffset); + } else { + const parsedQuery = query.getParsedQuery(); + if (!(parsedQuery instanceof SqlQuery)) return; + newQueryString = parsedQuery.apply(queryAction).toString(); + } + onQueryChange(query.changeQueryString(newQueryString)); + + if (shouldAutoRun()) { + setTimeout(() => { + if (typeof sliceIndex === 'number') { + const slice = findAllSqlQueriesInText(newQueryString)[sliceIndex]; + if (slice) { + void handleRun(false, slice); + } + } else { + void handleRun(false); + } + }, 20); + } + }, + ); function shouldAutoRun(): boolean { if (query.getEffectiveEngine() !== 'sql-native') return false; @@ -296,38 +322,37 @@ export const QueryTab = React.memo(function QueryTab(props: QueryTabProps) { if (querySlice) { effectiveQuery = effectiveQuery .changeQueryString(querySlice.sql) + .changeQueryContext({ ...effectiveQuery.queryContext, sliceIndex: querySlice.index }) .changePrefixLines(querySlice.startRowColumn.row); } - if (effectiveQuery.getEffectiveEngine() !== 'sql-msq-task') { - WorkbenchHistory.addQueryToHistory(effectiveQuery); - queryManager.runQuery(effectiveQuery); - return; - } - - effectiveQuery = preview - ? effectiveQuery.makePreview() - : effectiveQuery.setMaxNumTasksIfUnset(clusterCapacity); - - const capacityInfo = await maybeGetClusterCapacity(); - - const effectiveMaxNumTasks = effectiveQuery.queryContext.maxNumTasks ?? 2; - if (capacityInfo && capacityInfo.availableTaskSlots < effectiveMaxNumTasks) { - setAlertElement( - { - queryManager.runQuery(effectiveQuery); - }} - onClose={() => { - setAlertElement(undefined); - }} - />, - ); + if (effectiveQuery.getEffectiveEngine() === 'sql-msq-task') { + effectiveQuery = preview + ? effectiveQuery.makePreview() + : effectiveQuery.setMaxNumTasksIfUnset(clusterCapacity); + + const capacityInfo = await maybeGetClusterCapacity(); + + const effectiveMaxNumTasks = effectiveQuery.queryContext.maxNumTasks ?? 2; + if (capacityInfo && capacityInfo.availableTaskSlots < effectiveMaxNumTasks) { + setAlertElement( + { + queryManager.runQuery(effectiveQuery); + }} + onClose={() => { + setAlertElement(undefined); + }} + />, + ); + return; + } } else { - queryManager.runQuery(effectiveQuery); + WorkbenchHistory.addQueryToHistory(effectiveQuery); } + queryManager.runQuery(effectiveQuery); }); const statsTaskId: string | undefined = execution?.id; diff --git a/web-console/src/views/workbench-view/result-table-pane/result-table-pane.tsx b/web-console/src/views/workbench-view/result-table-pane/result-table-pane.tsx index cb38f6e0ca68..492767b3eaf8 100644 --- a/web-console/src/views/workbench-view/result-table-pane/result-table-pane.tsx +++ b/web-console/src/views/workbench-view/result-table-pane/result-table-pane.tsx @@ -44,6 +44,7 @@ import { columnToWidth, convertToGroupByExpression, copyAndAlert, + deepGet, filterMap, formatNumber, getNumericColumnBraces, @@ -80,7 +81,7 @@ function getExpressionIfAlias(query: SqlQuery, selectIndex: number): string { export interface ResultTablePaneProps { queryResult: QueryResult; - onQueryAction(action: QueryAction): void; + onQueryAction(action: QueryAction, sliceIndex?: number): void; onExport?(): void; runeMode: boolean; initPageSize?: number; @@ -103,6 +104,10 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result }); }, [queryResult.rows.length]); + function handleQueryAction(action: QueryAction) { + onQueryAction(action, deepGet(queryResult, 'query.context.sliceIndex')); + } + function hasFilterOnHeader(header: string, headerIndex: number): boolean { if (!parsedQuery || !parsedQuery.isRealOutputColumnAtSelectIndex(headerIndex)) return false; @@ -139,7 +144,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={reverseOrderByDirection === 'ASC' ? IconNames.SORT_ASC : IconNames.SORT_DESC} text={`Order ${reverseOrderByDirection === 'ASC' ? 'ascending' : 'descending'}`} onClick={() => { - onQueryAction(q => q.changeOrderByExpressions([reverseOrderBy])); + handleQueryAction(q => q.changeOrderByExpressions([reverseOrderBy])); }} />, ); @@ -150,7 +155,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.SORT_DESC} text="Order descending" onClick={() => { - onQueryAction(q => q.changeOrderByExpressions([descOrderBy])); + handleQueryAction(q => q.changeOrderByExpressions([descOrderBy])); }} />, { - onQueryAction(q => q.changeOrderByExpressions([ascOrderBy])); + handleQueryAction(q => q.changeOrderByExpressions([ascOrderBy])); }} />, ); @@ -175,7 +180,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result text="Remove cast" onClick={() => { if (!selectExpression || !underlyingExpression) return; - onQueryAction(q => + handleQueryAction(q => q.changeSelect( headerIndex, underlyingExpression.getArg(0)!.as(selectExpression.getOutputName()), @@ -196,7 +201,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result text={asType} onClick={() => { if (!selectExpression) return; - onQueryAction(q => + handleQueryAction(q => q.changeSelect( headerIndex, selectExpression @@ -238,7 +243,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result text={path} onClick={() => { if (!selectExpression) return; - onQueryAction(q => + handleQueryAction(q => q.addSelect( F('JSON_VALUE', selectExpression.getUnderlyingExpression(), path).as( selectExpression.getOutputName() + path.replace(/^\$/, ''), @@ -264,7 +269,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.FILTER_REMOVE} text="Remove from WHERE clause" onClick={() => { - onQueryAction(q => + handleQueryAction(q => q.changeWhereExpression(whereExpression.removeColumnFromAnd(header)), ); }} @@ -280,7 +285,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.FILTER_REMOVE} text="Remove from HAVING clause" onClick={() => { - onQueryAction(q => + handleQueryAction(q => q.changeHavingExpression(havingExpression.removeColumnFromAnd(header)), ); }} @@ -309,7 +314,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result key="time_floor" expression={selectExpression} onChange={expression => { - onQueryAction(q => q.changeSelect(headerIndex, expression)); + handleQueryAction(q => q.changeSelect(headerIndex, expression)); }} />, ); @@ -320,7 +325,9 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.TIME} text="Use as the primary time column" onClick={() => { - onQueryAction(q => q.changeSelect(headerIndex, selectExpression.as(TIME_COLUMN))); + handleQueryAction(q => + q.changeSelect(headerIndex, selectExpression.as(TIME_COLUMN)), + ); }} />, ); @@ -341,7 +348,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.TIME} text={`Time parse as '${possibleDruidFormat}' and use as the primary time column`} onClick={() => { - onQueryAction(q => + handleQueryAction(q => q.changeSelect(headerIndex, newSelectExpression.as(TIME_COLUMN)), ); }} @@ -352,7 +359,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result if (parsedQuery.hasGroupBy()) { if (parsedQuery.isGroupedOutputColumn(header)) { const convertToAggregate = (aggregate: SqlExpression) => { - onQueryAction(q => + handleQueryAction(q => q.removeOutputColumn(header).addSelect(aggregate, { insertIndex: 'last', }), @@ -428,7 +435,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.EXCHANGE} text="Convert to group by" onClick={() => { - onQueryAction(q => + handleQueryAction(q => q.removeOutputColumn(header).addSelect(groupByExpression, { insertIndex: 'last-grouping', addToGroupBy: 'end', @@ -450,7 +457,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result icon={IconNames.CROSS} text="Remove column" onClick={() => { - onQueryAction(q => q.removeOutputColumn(header)); + handleQueryAction(q => q.removeOutputColumn(header)); }} />, ); @@ -506,7 +513,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result headerIndex={headerIndex} runeMode={runeMode} query={parsedQuery} - onQueryAction={onQueryAction} + onQueryAction={handleQueryAction} onShowFullValue={setShowValue} /> ); @@ -640,7 +647,7 @@ export const ResultTablePane = React.memo(function ResultTablePane(props: Result expression={editingExpression} onSave={newExpression => { if (!parsedQuery) return; - onQueryAction(q => q.changeSelect(editingColumn, newExpression)); + handleQueryAction(q => q.changeSelect(editingColumn, newExpression)); }} onClose={() => setEditingColumn(-1)} />