From 9bd5237ed787e5a81e4299ed98df547dbb1e04fe Mon Sep 17 00:00:00 2001 From: Spencer Torres Date: Tue, 1 Oct 2024 07:51:22 -0400 Subject: [PATCH] Fix AST parsing when switching back to query builder (#1002) Co-authored-by: Andreas Christou --- CHANGELOG.md | 10 ++ .../queryBuilder/EditorTypeSwitcher.tsx | 18 +-- src/components/queryBuilder/utils.test.ts | 120 ++++++++++++++++++ src/components/queryBuilder/utils.ts | 29 +++-- src/data/utils.test.ts | 70 +++++++++- src/data/utils.ts | 32 ++++- src/views/CHQueryEditor.tsx | 4 +- 7 files changed, 262 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 89607cea..a60744c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## Unreleased + +### Features + +- Queries parsed from the SQL editor will now attempt to re-map columns into their correct fields for Log and Trace queries. + +### Fixes + +- Fixed and enhanced the logic for parsing a query back into the query builder. + ## 4.4.0 ### Features diff --git a/src/components/queryBuilder/EditorTypeSwitcher.tsx b/src/components/queryBuilder/EditorTypeSwitcher.tsx index abc14c19..625be5c3 100644 --- a/src/components/queryBuilder/EditorTypeSwitcher.tsx +++ b/src/components/queryBuilder/EditorTypeSwitcher.tsx @@ -6,13 +6,14 @@ import { generateSql } from 'data/sqlGenerator'; import labels from 'labels'; import { EditorType, CHQuery, defaultCHBuilderQuery } from 'types/sql'; import { QueryBuilderOptions } from 'types/queryBuilder'; -import isString from 'lodash/isString'; import { mapQueryTypeToGrafanaFormat } from 'data/utils'; +import { Datasource } from 'data/CHDatasource'; interface CHEditorTypeSwitcherProps { query: CHQuery; onChange: (query: CHQuery) => void; onRunQuery: () => void; + datasource?: Datasource; } const options: Array> = [ @@ -24,7 +25,7 @@ const options: Array> = [ * Component for switching between the SQL and Query Builder editors. */ export const EditorTypeSwitcher = (props: CHEditorTypeSwitcherProps) => { - const { query, onChange } = props; + const { datasource, query, onChange } = props; const { label, tooltip, switcher, cannotConvert } = labels.components.EditorTypeSwitcher; const editorType: EditorType = query.editorType || EditorType.Builder; const [confirmModalState, setConfirmModalState] = useState(false); @@ -33,12 +34,12 @@ export const EditorTypeSwitcher = (props: CHEditorTypeSwitcherProps) => { const onEditorTypeChange = (editorType: EditorType, confirmed = false) => { // TODO: component state has updated, but not local state. if (query.editorType === EditorType.SQL && editorType === EditorType.Builder && !confirmed) { - const queryOptionsFromSql = getQueryOptionsFromSql(query.rawSql); - if (isString(queryOptionsFromSql)) { - setCannotConvertModalState(true); - setErrorMessage(queryOptionsFromSql); - } else { + try { + getQueryOptionsFromSql(query.rawSql, query.queryType, datasource); setConfirmModalState(true); + } catch (err) { + setCannotConvertModalState(true); + setErrorMessage((err as Error).message); } } else { let builderOptions: QueryBuilderOptions; @@ -47,7 +48,7 @@ export const EditorTypeSwitcher = (props: CHEditorTypeSwitcherProps) => { builderOptions = query.builderOptions; break; case EditorType.SQL: - builderOptions = getQueryOptionsFromSql(query.rawSql) as QueryBuilderOptions; + builderOptions = getQueryOptionsFromSql(query.rawSql, query.queryType, datasource) as QueryBuilderOptions; break; default: builderOptions = defaultCHBuilderQuery.builderOptions; @@ -66,6 +67,7 @@ export const EditorTypeSwitcher = (props: CHEditorTypeSwitcherProps) => { onChange({ ...query, editorType: EditorType.Builder, + queryType: builderOptions.queryType, rawSql: generateSql(builderOptions), builderOptions }); diff --git a/src/components/queryBuilder/utils.test.ts b/src/components/queryBuilder/utils.test.ts index e29fe04a..193c1d45 100644 --- a/src/components/queryBuilder/utils.test.ts +++ b/src/components/queryBuilder/utils.test.ts @@ -1,6 +1,8 @@ import { generateSql } from 'data/sqlGenerator'; import { getQueryOptionsFromSql, isDateTimeType, isDateType, isNumberType } from './utils'; import { AggregateType, BuilderMode, ColumnHint, DateFilterWithoutValue, FilterOperator, MultiFilter, OrderByDirection, QueryBuilderOptions, QueryType } from 'types/queryBuilder'; +import { Datasource } from 'data/CHDatasource'; +import otel from 'otel'; describe('isDateType', () => { it('returns true for Date type', () => { @@ -450,6 +452,124 @@ describe('getQueryOptionsFromSql', () => { false ); + testCondition('handles parsing a column with a complex name with spaces and capital characters', 'SELECT "Complex Name" FROM "db"."foo"', { + queryType: QueryType.Table, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'Complex Name', alias: undefined }], + aggregates: [], + }); + + it('matches input query type', () => { + const sql = 'SELECT test FROM "db"."foo"'; + const expectedOptions: QueryBuilderOptions = { + queryType: QueryType.Logs, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'test', alias: undefined }], + aggregates: [], + }; + + expect(getQueryOptionsFromSql(sql, QueryType.Logs)).toEqual(expectedOptions); + }); + + it('matches column hints with Grafana query aliases', () => { + const sql = 'SELECT a as body, b as level FROM "db"."foo"'; + const expectedOptions: QueryBuilderOptions = { + queryType: QueryType.Logs, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'a', alias: 'body', hint: ColumnHint.LogMessage }, { name: 'b', alias: 'level', hint: ColumnHint.LogLevel }], + aggregates: [], + limit: undefined + }; + + expect(getQueryOptionsFromSql(sql, QueryType.Logs)).toEqual(expectedOptions); + }); + + it('matches column hints with OTel log column names', () => { + const mockDs = {} as Datasource; + mockDs.getDefaultLogsColumns = jest.fn(() => otel.getLatestVersion().logColumnMap); + + const sql = 'SELECT "Timestamp", "SeverityText" FROM "db"."foo"'; + const expectedOptions: QueryBuilderOptions = { + queryType: QueryType.Logs, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'Timestamp', alias: undefined, hint: ColumnHint.Time }, { name: 'SeverityText', alias: undefined, hint: ColumnHint.LogLevel }], + aggregates: [], + limit: undefined + }; + + expect(getQueryOptionsFromSql(sql, QueryType.Logs, mockDs)).toEqual(expectedOptions); + }); + + it('matches column hints with datasource log column names', () => { + const mockDs = {} as Datasource; + mockDs.getDefaultLogsColumns = jest.fn(() => ( + new Map([ + [ColumnHint.Time, 'SpecialTimestamp'], + [ColumnHint.LogMessage, 'LogBody'] + ]))); + + const sql = 'SELECT "SpecialTimestamp", "LogBody" FROM "db"."foo"'; + const expectedOptions: QueryBuilderOptions = { + queryType: QueryType.Logs, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'SpecialTimestamp', alias: undefined, hint: ColumnHint.Time }, { name: 'LogBody', alias: undefined, hint: ColumnHint.LogMessage }], + aggregates: [], + limit: undefined + }; + + expect(getQueryOptionsFromSql(sql, QueryType.Logs, mockDs)).toEqual(expectedOptions); + }); + + it('matches column hints with OTel trace column names', () => { + const mockDs = {} as Datasource; + mockDs.getDefaultTraceColumns = jest.fn(() => otel.getLatestVersion().traceColumnMap); + + const sql = 'SELECT "StartTime", "ServiceName" FROM "db"."foo"'; + const expectedOptions: QueryBuilderOptions = { + queryType: QueryType.Traces, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'StartTime', alias: undefined, hint: ColumnHint.Time }, { name: 'ServiceName', alias: undefined, hint: ColumnHint.TraceServiceName }], + aggregates: [], + limit: undefined + }; + + expect(getQueryOptionsFromSql(sql, QueryType.Traces, mockDs)).toEqual(expectedOptions); + }); + + it('matches column hints with datasource trace column names', () => { + const mockDs = {} as Datasource; + mockDs.getDefaultTraceColumns = jest.fn(() => ( + new Map([ + [ColumnHint.Time, 'SpecialTimestamp'], + [ColumnHint.TraceId, 'CustomTraceID'] + ]))); + + const sql = 'SELECT "SpecialTimestamp", "CustomTraceID" FROM "db"."foo"'; + const expectedOptions: QueryBuilderOptions = { + queryType: QueryType.Traces, + mode: BuilderMode.List, + database: 'db', + table: 'foo', + columns: [{ name: 'SpecialTimestamp', alias: undefined, hint: ColumnHint.Time }, { name: 'CustomTraceID', alias: undefined, hint: ColumnHint.TraceId }], + aggregates: [], + limit: undefined + }; + + expect(getQueryOptionsFromSql(sql, QueryType.Traces, mockDs)).toEqual(expectedOptions); + }); + it('Handles brackets and Grafana macros/variables', () => { const sql = ` /* \${__variable} \${__variable.key} */ diff --git a/src/components/queryBuilder/utils.ts b/src/components/queryBuilder/utils.ts index a1cf5677..1259c1ef 100644 --- a/src/components/queryBuilder/utils.ts +++ b/src/components/queryBuilder/utils.ts @@ -32,7 +32,9 @@ import { QueryType, } from 'types/queryBuilder'; import { sqlToStatement } from 'data/ast'; -import { getColumnByHint } from 'data/sqlGenerator'; +import { getColumnByHint, logColumnHintsToAlias } from 'data/sqlGenerator'; +import { Datasource } from 'data/CHDatasource'; +import { tryApplyColumnHints } from 'data/utils'; export const isBooleanType = (type: string): boolean => { @@ -79,19 +81,19 @@ export const isMultiFilter = (filter: Filter): filter is MultiFilter => { return isStringType(filter.type) && [FilterOperator.In, FilterOperator.NotIn].includes(filter.operator); }; -export function getQueryOptionsFromSql(sql: string): QueryBuilderOptions | string { +export function getQueryOptionsFromSql(sql: string, queryType?: QueryType, datasource?: Datasource): QueryBuilderOptions { const ast = sqlToStatement(sql); if (!ast) { - return 'The query is not valid SQL.'; + throw new Error('The query is not valid SQL.'); } if (ast.type !== 'select') { - return 'The query is not a select statement.'; + throw new Error('The query is not a select statement.'); } if (!ast.from || ast.from.length !== 1) { - return `The query has too many 'FROM' clauses.`; + throw new Error(`The query has too many 'FROM' clauses.`); } if (ast.from[0].type !== 'table') { - return `The 'FROM' clause is not a table.`; + throw new Error(`The 'FROM' clause is not a table.`); } const fromTable = ast.from[0] as FromTable; @@ -100,14 +102,22 @@ export function getQueryOptionsFromSql(sql: string): QueryBuilderOptions | strin const builderOptions = { database: fromTable.name.schema || '', table: fromTable.name.name || '', - queryType: QueryType.Table, + queryType: queryType || QueryType.Table, mode: BuilderMode.List, columns: [], aggregates: [], } as QueryBuilderOptions; if (columnsAndAggregates.columns.length > 0) { - builderOptions.columns = columnsAndAggregates.columns; + builderOptions.columns = columnsAndAggregates.columns || []; + } + + // Reconstruct column hints based off of known column names / aliases + if (queryType === QueryType.Logs) { + tryApplyColumnHints(builderOptions.columns!, datasource?.getDefaultLogsColumns()); // Try match default log columns + tryApplyColumnHints(builderOptions.columns!, logColumnHintsToAlias); // Try match Grafana aliases + } else if (queryType === QueryType.Traces) { + tryApplyColumnHints(builderOptions.columns!, datasource?.getDefaultTraceColumns()); } if (columnsAndAggregates.aggregates.length > 0) { @@ -116,7 +126,7 @@ export function getQueryOptionsFromSql(sql: string): QueryBuilderOptions | strin } const timeColumn = getColumnByHint(builderOptions, ColumnHint.Time); - if (timeColumn) { + if (!queryType && timeColumn) { builderOptions.queryType = QueryType.TimeSeries; if (builderOptions.aggregates?.length || 0) { builderOptions.mode = BuilderMode.Trend; @@ -153,6 +163,7 @@ export function getQueryOptionsFromSql(sql: string): QueryBuilderOptions | strin if (groupBy && groupBy.length > 0) { builderOptions.groupBy = groupBy; } + return builderOptions; } diff --git a/src/data/utils.test.ts b/src/data/utils.test.ts index c22dfda2..14de85a5 100644 --- a/src/data/utils.test.ts +++ b/src/data/utils.test.ts @@ -1,5 +1,5 @@ import { ColumnHint, QueryBuilderOptions, QueryType } from "types/queryBuilder"; -import { columnLabelToPlaceholder, dataFrameHasLogLabelWithName, isBuilderOptionsRunnable, transformQueryResponseWithTraceAndLogLinks } from "./utils"; +import { columnLabelToPlaceholder, dataFrameHasLogLabelWithName, isBuilderOptionsRunnable, transformQueryResponseWithTraceAndLogLinks, tryApplyColumnHints } from "./utils"; import { newMockDatasource } from "__mocks__/datasource"; import { CoreApp, DataFrame, DataQueryRequest, DataQueryResponse, Field, FieldType } from "@grafana/data"; import { CHBuilderQuery, CHQuery, EditorType } from "types/sql"; @@ -32,6 +32,74 @@ describe('isBuilderOptionsRunnable', () => { }); }); +describe('tryApplyColumnHints', () => { + it('does not apply hints when queryType and hint map are not provided', () => { + const columns = [ + { name: 'a', alias: undefined, hint: undefined }, + { name: 'b', alias: undefined, hint: undefined }, + ]; + + tryApplyColumnHints(columns); + + expect(columns[0].hint).toBeUndefined(); + expect(columns[1].hint).toBeUndefined(); + }); + + it('applies time hint to columns that contain "time"', () => { + const columns = [ + { name: 'Timestamp', alias: undefined, hint: undefined }, + { name: 'log_timestamp', alias: undefined, hint: undefined }, + ]; + + tryApplyColumnHints(columns); + + expect(columns[0].hint).toEqual(ColumnHint.Time); + expect(columns[1].hint).toEqual(ColumnHint.Time); + }); + + it('does not apply hints to column with existing hint', () => { + const columns = [ + { name: 'time', alias: undefined, hint: ColumnHint.TraceServiceName }, + ]; + + tryApplyColumnHints(columns); + + expect(columns[0].hint).toEqual(ColumnHint.TraceServiceName); + }); + + it('applies hints by column name according to hint map, ignoring case', () => { + const columns = [ + { name: 'Super_Custom_Timestamp', alias: undefined, hint: undefined }, + { name: 'LogLevel', alias: undefined, hint: undefined }, + ]; + const hintMap: Map = new Map([ + [ColumnHint.Time, "super_custom_timestamp"], + [ColumnHint.LogLevel, "LogLevel"] + ]); + + tryApplyColumnHints(columns, hintMap); + + expect(columns[0].hint).toEqual(ColumnHint.Time); + expect(columns[1].hint).toEqual(ColumnHint.LogLevel); + }); + + it('applies hints by column alias according to hint map, ignoring case', () => { + const columns = [ + { name: 'other name', alias: 'Super_Custom_Timestamp', hint: undefined }, + { name: 'other name', alias: 'LogLevel', hint: undefined }, + ]; + const hintMap: Map = new Map([ + [ColumnHint.Time, "super_custom_timestamp"], + [ColumnHint.LogLevel, "LogLevel"] + ]); + + tryApplyColumnHints(columns, hintMap); + + expect(columns[0].hint).toEqual(ColumnHint.Time); + expect(columns[1].hint).toEqual(ColumnHint.LogLevel); + }); +}); + describe('columnLabelToPlaceholder', () => { it('converts to lowercase and removes multiple spaces', () => { const expected = 'expected_test_output'; diff --git a/src/data/utils.ts b/src/data/utils.ts index 492e7207..bd8f68ca 100644 --- a/src/data/utils.ts +++ b/src/data/utils.ts @@ -1,5 +1,5 @@ import { CoreApp, DataFrame, DataQueryRequest, DataQueryResponse } from "@grafana/data"; -import { ColumnHint, FilterOperator, OrderByDirection, QueryBuilderOptions, QueryType, StringFilter } from "types/queryBuilder" +import { ColumnHint, FilterOperator, OrderByDirection, QueryBuilderOptions, QueryType, SelectedColumn, StringFilter } from "types/queryBuilder" import { CHBuilderQuery, CHQuery, EditorType } from "types/sql"; import { Datasource } from "./CHDatasource"; import { pluginVersion } from "utils/version"; @@ -75,6 +75,36 @@ export const mapGrafanaFormatToQueryType = (f?: number): QueryType => { } }; +/** + * Manipulates column array in-place to include column hints, loosely matched by the provided column hint map. + */ +export const tryApplyColumnHints = (columns: SelectedColumn[], hintsToColumns?: Map) => { + const columnsToHints: Map = new Map(); + if (hintsToColumns) { + hintsToColumns.forEach((name, hint) => { + columnsToHints.set(name.toLowerCase().trim(), hint); + }); + } + + for (const column of columns) { + if (column.hint) { + continue; + } + + const name = column.name.toLowerCase().trim(); + const alias = column.alias?.toLowerCase().trim() || ''; + + const hint = columnsToHints.get(name) || columnsToHints.get(alias); + if (hint) { + column.hint = hint; + continue; + } + + if (name.includes('time')) { + column.hint = ColumnHint.Time; + } + } +}; /** * Converts label into sql-style column name. diff --git a/src/views/CHQueryEditor.tsx b/src/views/CHQueryEditor.tsx index 35b1deb5..6561e42c 100644 --- a/src/views/CHQueryEditor.tsx +++ b/src/views/CHQueryEditor.tsx @@ -20,13 +20,13 @@ export type CHQueryEditorProps = QueryEditorProps * Top level query editor component */ export const CHQueryEditor = (props: CHQueryEditorProps) => { - const { query: savedQuery, onRunQuery } = props; + const { datasource, query: savedQuery, onRunQuery } = props; const query = migrateCHQuery(savedQuery); return ( <>
- +