Fix AST parsing when switching back to query builder

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

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<SelectableValue<EditorType>> = [
@@ -24,7 +25,7 @@ const options: Array<SelectableValue<EditorType>> = [
  * 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<boolean>(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
         });

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} */

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

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<ColumnHint, string> = 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<ColumnHint, string> = 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';