[Lens] Escape field names in formula (#102588)

* [Lens] Escape field names in formula

* Fix handling of partially typed fields with invalid chars

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

packages/kbn-tinymath/grammar/grammar.peggy

@@ -43,7 +43,7 @@ Literal "literal"
 // Quoted variables are interpreted as strings
 // but unquoted variables are more restrictive
 Variable
-  = _ [\'] chars:(ValidChar / Space / [\"])* [\'] _ {
+  = _ '"' chars:("\\\"" { return "\""; } / [^"])* '"' _ {
     return {
       type: 'variable',
       value: chars.join(''),
@@ -51,7 +51,7 @@ Variable
       text: text()
     };
   }
-  / _ [\"] chars:(ValidChar / Space / [\'])* [\"] _ {
+  / _ "'" chars:("\\\'" { return "\'"; } / [^'])* "'" _ {
       return {
       type: 'variable',
       value: chars.join(''),

packages/kbn-tinymath/test/library.test.js

@@ -92,6 +92,7 @@ describe('Parser', () => {
       expect(parse('@foo0')).toEqual(variableEqual('@foo0'));
       expect(parse('.foo0')).toEqual(variableEqual('.foo0'));
       expect(parse('-foo0')).toEqual(variableEqual('-foo0'));
+      expect(() => parse(`foo😀\t')`)).toThrow('Failed to parse');
     });
   });
 
@@ -103,6 +104,7 @@ describe('Parser', () => {
       expect(parse('"foo bar fizz buzz"')).toEqual(variableEqual('foo bar fizz buzz'));
       expect(parse('"foo   bar   baby"')).toEqual(variableEqual('foo   bar   baby'));
       expect(parse(`"f'oo"`)).toEqual(variableEqual(`f'oo`));
+      expect(parse(`"foo😀\t"`)).toEqual(variableEqual(`foo😀\t`));
     });
 
     it('strings with single quotes', () => {
@@ -119,6 +121,7 @@ describe('Parser', () => {
       expect(parse("'foo bar '")).toEqual(variableEqual("foo bar "));
       expect(parse("'0foo'")).toEqual(variableEqual("0foo"));
       expect(parse(`'f"oo'`)).toEqual(variableEqual(`f"oo`));
+      expect(parse(`'foo😀\t'`)).toEqual(variableEqual(`foo😀\t`));
       /* eslint-enable prettier/prettier */
     });
 

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/editor/formula_editor.tsx

@@ -29,7 +29,7 @@ import { ParamEditorProps } from '../../index';
 import { getManagedColumnsFrom } from '../../../layer_helpers';
 import { ErrorWrapper, runASTValidation, tryToParse } from '../validation';
 import {
-  LensMathSuggestion,
+  LensMathSuggestions,
   SUGGESTION_TYPE,
   suggest,
   getSuggestion,
@@ -329,7 +329,7 @@ export function FormulaEditor({
       context: monaco.languages.CompletionContext
     ) => {
       const innerText = model.getValue();
-      let aSuggestions: { list: LensMathSuggestion[]; type: SUGGESTION_TYPE } = {
+      let aSuggestions: LensMathSuggestions = {
         list: [],
         type: SUGGESTION_TYPE.FIELD,
       };
@@ -367,7 +367,13 @@ export function FormulaEditor({
 
       return {
         suggestions: aSuggestions.list.map((s) =>
-          getSuggestion(s, aSuggestions.type, visibleOperationsMap, context.triggerCharacter)
+          getSuggestion(
+            s,
+            aSuggestions.type,
+            visibleOperationsMap,
+            context.triggerCharacter,
+            aSuggestions.range
+          )
         ),
       };
     },

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/editor/math_completion.test.ts

@@ -18,6 +18,7 @@ import {
   getHover,
   suggest,
   monacoPositionToOffset,
+  offsetToRowColumn,
   getInfoAtZeroIndexedPosition,
 } from './math_completion';
 
@@ -363,6 +364,36 @@ describe('math completion', () => {
     });
   });
 
+  describe('offsetToRowColumn', () => {
+    it('should work with single-line strings', () => {
+      const input = `0123456`;
+      expect(offsetToRowColumn(input, 5)).toEqual(
+        expect.objectContaining({
+          lineNumber: 1,
+          column: 6,
+        })
+      );
+    });
+
+    it('should work with multi-line strings accounting for newline characters', () => {
+      const input = `012
+456
+89')`;
+      expect(offsetToRowColumn(input, 0)).toEqual(
+        expect.objectContaining({
+          lineNumber: 1,
+          column: 1,
+        })
+      );
+      expect(offsetToRowColumn(input, 9)).toEqual(
+        expect.objectContaining({
+          lineNumber: 3,
+          column: 2,
+        })
+      );
+    });
+  });
+
   describe('monacoPositionToOffset', () => {
     it('should work with multi-line strings accounting for newline characters', () => {
       const input = `012

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/editor/math_completion.ts

@@ -13,6 +13,7 @@ import {
   TinymathLocation,
   TinymathAST,
   TinymathFunction,
+  TinymathVariable,
   TinymathNamedArgument,
 } from '@kbn/tinymath';
 import type {
@@ -21,7 +22,7 @@ import type {
 } from '../../../../../../../../../src/plugins/data/public';
 import { IndexPattern } from '../../../../types';
 import { memoizedGetAvailableOperationsByMetadata } from '../../../operations';
-import { tinymathFunctions, groupArgsByType } from '../util';
+import { tinymathFunctions, groupArgsByType, unquotedStringRegex } from '../util';
 import type { GenericOperationDefinition } from '../..';
 import { getFunctionSignatureLabel, getHelpTextContent } from './formula_help';
 import { hasFunctionFieldArgument } from '../validation';
@@ -47,6 +48,7 @@ export type LensMathSuggestion =
 export interface LensMathSuggestions {
   list: LensMathSuggestion[];
   type: SUGGESTION_TYPE;
+  range?: monaco.IRange;
 }
 
 function inLocation(cursorPosition: number, location: TinymathLocation) {
@@ -92,7 +94,7 @@ export function offsetToRowColumn(expression: string, offset: number): monaco.Po
   let lineNumber = 1;
   for (const line of lines) {
     if (line.length >= remainingChars) {
-      return new monaco.Position(lineNumber, remainingChars);
+      return new monaco.Position(lineNumber, remainingChars + 1);
     }
     remainingChars -= line.length + 1;
     lineNumber++;
@@ -128,7 +130,7 @@ export async function suggest({
   operationDefinitionMap: Record<string, GenericOperationDefinition>;
   data: DataPublicPluginStart;
   dateHistogramInterval?: number;
-}): Promise<{ list: LensMathSuggestion[]; type: SUGGESTION_TYPE }> {
+}): Promise<LensMathSuggestions> {
   const text =
     expression.substr(0, zeroIndexedOffset) + MARKER + expression.substr(zeroIndexedOffset);
   try {
@@ -154,6 +156,7 @@ export async function suggest({
       return getArgumentSuggestions(
         tokenInfo.parent,
         tokenInfo.parent.args.findIndex((a) => a === tokenAst),
+        text,
         indexPattern,
         operationDefinitionMap
       );
@@ -210,6 +213,7 @@ function getFunctionSuggestions(
 function getArgumentSuggestions(
   ast: TinymathFunction,
   position: number,
+  expression: string,
   indexPattern: IndexPattern,
   operationDefinitionMap: Record<string, GenericOperationDefinition>
 ) {
@@ -280,7 +284,16 @@ function getArgumentSuggestions(
         .filter((op) => op.operationType === operation.type)
         .map((op) => ('field' in op ? op.field : undefined))
         .filter((field) => field);
-      return { list: fields as string[], type: SUGGESTION_TYPE.FIELD };
+      const fieldArg = ast.args[0];
+      const location = typeof fieldArg !== 'string' && (fieldArg as TinymathVariable).location;
+      let range: monaco.IRange | undefined;
+      if (location) {
+        const start = offsetToRowColumn(expression, location.min);
+        // This accounts for any characters that the user has already typed
+        const end = offsetToRowColumn(expression, location.max - MARKER.length);
+        range = monaco.Range.fromPositions(start, end);
+      }
+      return { list: fields as string[], type: SUGGESTION_TYPE.FIELD, range };
     } else {
       return { list: [], type: SUGGESTION_TYPE.FIELD };
     }
@@ -375,7 +388,8 @@ export function getSuggestion(
   suggestion: LensMathSuggestion,
   type: SUGGESTION_TYPE,
   operationDefinitionMap: Record<string, GenericOperationDefinition>,
-  triggerChar: string | undefined
+  triggerChar: string | undefined,
+  range?: monaco.IRange
 ): monaco.languages.CompletionItem {
   let kind: monaco.languages.CompletionItemKind = monaco.languages.CompletionItemKind.Method;
   let label: string =
@@ -397,6 +411,10 @@ export function getSuggestion(
       break;
     case SUGGESTION_TYPE.FIELD:
       kind = monaco.languages.CompletionItemKind.Value;
+      // Look for unsafe characters
+      if (unquotedStringRegex.test(label)) {
+        insertText = `'${label.replaceAll(`'`, "\\'")}'`;
+      }
       break;
     case SUGGESTION_TYPE.FUNCTIONS:
       insertText = `${label}($0)`;
@@ -450,7 +468,7 @@ export function getSuggestion(
     command,
     additionalTextEdits: [],
     // @ts-expect-error Monaco says this type is required, but provides a default value
-    range: undefined,
+    range,
     sortText,
     filterText,
   };

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/generate.ts

@@ -13,6 +13,7 @@ import {
 } from '../index';
 import { ReferenceBasedIndexPatternColumn } from '../column_types';
 import { IndexPatternLayer } from '../../../types';
+import { unquotedStringRegex } from './util';
 
 // Just handle two levels for now
 type OperationParams = Record<string, string | number | Record<string, string | number>>;
@@ -25,6 +26,9 @@ export function getSafeFieldName({
   if (!fieldName || operationType === 'count') {
     return '';
   }
+  if (unquotedStringRegex.test(fieldName)) {
+    return `'${fieldName.replaceAll(`'`, "\\'")}'`;
+  }
   return fieldName;
 }
 

x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/util.ts

@@ -16,6 +16,8 @@ import type {
 import type { OperationDefinition, IndexPatternColumn, GenericOperationDefinition } from '../index';
 import type { GroupedNodes } from './types';
 
+export const unquotedStringRegex = /[^0-9A-Za-z._@\[\]/]/;
+
 export function groupArgsByType(args: TinymathAST[]) {
   const { namedArgument, variable, function: functions } = groupBy<TinymathAST>(
     args,

x-pack/test/functional/apps/lens/formula.ts

@@ -14,6 +14,7 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
   const listingTable = getService('listingTable');
   const browser = getService('browser');
   const testSubjects = getService('testSubjects');
+  const fieldEditor = getService('fieldEditor');
 
   describe('lens formula', () => {
     it('should transition from count to formula', async () => {
@@ -88,6 +89,41 @@ export default function ({ getService, getPageObjects }: FtrProviderContext) {
       expect(await element.getVisibleText()).to.equal(`count(kql='Men\\'s Clothing')`);
     });
 
+    it('should insert single quotes and escape when needed to create valid field name', async () => {
+      await PageObjects.visualize.navigateToNewVisualization();
+      await PageObjects.visualize.clickVisType('lens');
+      await PageObjects.lens.goToTimeRange();
+      await PageObjects.lens.switchToVisualization('lnsDatatable');
+      await PageObjects.lens.clickAddField();
+      await fieldEditor.setName(`*' "'`);
+      await fieldEditor.enableValue();
+      await fieldEditor.typeScript("emit('abc')");
+      await fieldEditor.save();
+
+      await PageObjects.lens.configureDimension({
+        dimension: 'lnsDatatable_metrics > lns-empty-dimension',
+        operation: 'unique_count',
+        field: `*`,
+        keepOpen: true,
+      });
+
+      await PageObjects.lens.switchToFormula();
+      let element = await find.byCssSelector('.monaco-editor');
+      expect(await element.getVisibleText()).to.equal(`unique_count('*\\' "\\'')`);
+
+      const input = await find.activeElement();
+      await input.clearValueWithKeyboard({ charByChar: true });
+      await input.type('unique_count(');
+      await PageObjects.common.sleep(100);
+      await input.type('*');
+      await input.pressKeys(browser.keys.ENTER);
+
+      await PageObjects.common.sleep(100);
+
+      element = await find.byCssSelector('.monaco-editor');
+      expect(await element.getVisibleText()).to.equal(`unique_count('*\\' "\\'')`);
+    });
+
     it('should persist a broken formula on close', async () => {
       await PageObjects.visualize.navigateToNewVisualization();
       await PageObjects.visualize.clickVisType('lens');