[Lens] Provide formula helper to simplify integration of Lens instances

Closes: elastic#103055
alexwizp · Jan 6, 2022 · 98aa9b3 · 98aa9b3
1 parent e62768a
commit 98aa9b3
Show file tree

Hide file tree

Showing 10 changed files with 264 additions and 140 deletions.
diff --git a/x-pack/plugins/lens/public/indexpattern_datasource/index.ts b/x-pack/plugins/lens/public/indexpattern_datasource/index.ts
@@ -7,6 +7,7 @@
 
 import type { CoreSetup } from 'kibana/public';
 import { createStartServicesGetter, Storage } from '../../../../../src/plugins/kibana_utils/public';
+import { expandFormulaColumn, generateFormulaColumns } from './operations/definitions/formula';
 import type { ExpressionsSetup } from '../../../../../src/plugins/expressions/public';
 import type { ChartsPluginSetup } from '../../../../../src/plugins/charts/public';
 import type { IndexPatternFieldEditorStart } from '../../../../../src/plugins/data_view_field_editor/public';
@@ -21,6 +22,11 @@ import type {
   FieldFormatsSetup,
 } from '../../../../../src/plugins/field_formats/public';
 
+export { expandFormulaColumn, generateFormulaColumns };
+
+export type ExpandFormulaColumnFn = typeof expandFormulaColumn;
+export type GenerateFormulaColumnsFn = typeof generateFormulaColumns;
+
 export interface IndexPatternDatasourceSetupPlugins {
   expressions: ExpressionsSetup;
   fieldFormats: FieldFormatsSetup;

diff --git a/...s/public/indexpattern_datasource/operations/definitions/formula/editor/formula_editor.tsx b/...s/public/indexpattern_datasource/operations/definitions/formula/editor/formula_editor.tsx
@@ -45,7 +45,7 @@ import { trackUiEvent } from '../../../../../lens_ui_telemetry';
 
 import './formula.scss';
 import { FormulaIndexPatternColumn } from '../formula';
-import { regenerateLayerFromAst } from '../parse';
+import { generateFormulaLayer, expandFormulaColumn } from '../parse';
 import { filterByVisibleOperation } from '../util';
 import { getColumnTimeShiftWarnings, getDateHistogramInterval } from '../../../../time_shift_utils';
 
@@ -151,16 +151,21 @@ export function FormulaEditor({
     setIsCloseable(true);
     // If the text is not synced, update the column.
     if (text !== currentColumn.params.formula) {
-      updateLayer((prevLayer) => {
-        return regenerateLayerFromAst(
-          text || '',
-          prevLayer,
-          columnId,
-          currentColumn,
+      updateLayer((prevLayer) =>
+        generateFormulaLayer({
+          id: columnId,
+          formula: text || '',
+          layer: {
+            ...prevLayer,
+            columns: {
+              ...prevLayer.columns,
+              [columnId]: { ...currentColumn },
+            },
+          },
           indexPattern,
-          operationDefinitionMap
-        ).newLayer;
-      });
+          operations: operationDefinitionMap,
+        })
+      );
     }
   });
 
@@ -173,15 +178,21 @@ export function FormulaEditor({
         monaco.editor.setModelMarkers(editorModel.current, 'LENS', []);
         if (currentColumn.params.formula) {
           // Only submit if valid
-          const { newLayer } = regenerateLayerFromAst(
-            text || '',
-            layer,
-            columnId,
-            currentColumn,
-            indexPattern,
-            operationDefinitionMap
+          updateLayer(
+            generateFormulaLayer({
+              id: columnId,
+              formula: text || '',
+              layer: {
+                ...layer,
+                columns: {
+                  ...layer.columns,
+                  [columnId]: { ...currentColumn },
+                },
+              },
+              indexPattern,
+              operations: operationDefinitionMap,
+            })
           );
-          updateLayer(newLayer);
         }
 
         return;
@@ -209,14 +220,19 @@ export function FormulaEditor({
           // If the formula is already broken, show the latest error message in the workspace
           if (currentColumn.params.formula !== text) {
             updateLayer(
-              regenerateLayerFromAst(
-                text || '',
-                layer,
-                columnId,
-                currentColumn,
+              generateFormulaLayer({
+                id: columnId,
+                formula: text || '',
+                layer: {
+                  ...layer,
+                  columns: {
+                    ...layer.columns,
+                    [columnId]: { ...currentColumn },
+                  },
+                },
                 indexPattern,
-                visibleOperationsMap
-              ).newLayer
+                operations: operationDefinitionMap,
+              })
             );
           }
         }
@@ -264,15 +280,29 @@ export function FormulaEditor({
         monaco.editor.setModelMarkers(editorModel.current, 'LENS', []);
 
         // Only submit if valid
-        const { newLayer, locations } = regenerateLayerFromAst(
-          text || '',
-          layer,
-          columnId,
-          currentColumn,
+        const {
+          columnOrder,
+          columns,
+          meta: { locations },
+        } = expandFormulaColumn({
+          id: columnId,
+          formula: text || '',
+          layer: {
+            ...layer,
+            columns: {
+              ...layer.columns,
+              [columnId]: currentColumn,
+            },
+          },
           indexPattern,
-          visibleOperationsMap
-        );
-        updateLayer(newLayer);
+          operations: operationDefinitionMap,
+        });
+
+        const newLayer = {
+          ...layer,
+          columnOrder,
+          columns,
+        };
 
         const managedColumns = getManagedColumnsFrom(columnId, newLayer.columns);
         const markers: monaco.editor.IMarkerData[] = managedColumns

diff --git a/...ugins/lens/public/indexpattern_datasource/operations/definitions/formula/formula.test.tsx b/...ugins/lens/public/indexpattern_datasource/operations/definitions/formula/formula.test.tsx
@@ -8,7 +8,7 @@
 import { createMockedIndexPattern } from '../../../mocks';
 import { formulaOperation, GenericOperationDefinition, GenericIndexPatternColumn } from '../index';
 import { FormulaIndexPatternColumn } from './formula';
-import { regenerateLayerFromAst } from './parse';
+import { expandFormulaColumn, generateFormulaLayer } from './parse';
 import type { IndexPattern, IndexPatternField, IndexPatternLayer } from '../../../types';
 import { tinymathFunctions } from './util';
 import { TermsIndexPatternColumn } from '../terms';
@@ -417,14 +417,19 @@ describe('formula', () => {
 
     function testIsBrokenFormula(formula: string) {
       expect(
-        regenerateLayerFromAst(
+        generateFormulaLayer({
           formula,
-          layer,
-          'col1',
-          currentColumn,
+          layer: {
+            ...layer,
+            columns: {
+              ...layer.columns,
+              col1: currentColumn,
+            },
+          },
+          id: 'col1',
           indexPattern,
-          operationDefinitionMap
-        ).newLayer
+          operations: operationDefinitionMap,
+        })
       ).toEqual({
         ...layer,
         columns: {
@@ -457,29 +462,23 @@ describe('formula', () => {
 
     it('should mutate the layer with new columns for valid formula expressions', () => {
       expect(
-        regenerateLayerFromAst(
-          'average(bytes)',
-          layer,
-          'col1',
-          currentColumn,
+        generateFormulaLayer({
+          formula: 'average(bytes)',
+          layer: {
+            ...layer,
+            columns: {
+              ...layer.columns,
+              col1: currentColumn,
+            },
+          },
+          id: 'col1',
           indexPattern,
-          operationDefinitionMap
-        ).newLayer
+          operations: operationDefinitionMap,
+        })
       ).toEqual({
         ...layer,
         columnOrder: ['col1X0', 'col1'],
         columns: {
-          ...layer.columns,
-          col1: {
-            ...currentColumn,
-            label: 'average(bytes)',
-            references: ['col1X0'],
-            params: {
-              ...currentColumn.params,
-              formula: 'average(bytes)',
-              isFormulaBroken: false,
-            },
-          },
           col1X0: {
             customLabel: true,
             dataType: 'number',
@@ -490,20 +489,35 @@ describe('formula', () => {
             sourceField: 'bytes',
             timeScale: false,
           },
+          col1: {
+            ...currentColumn,
+            label: 'average(bytes)',
+            references: ['col1X0'],
+            params: {
+              ...currentColumn.params,
+              formula: 'average(bytes)',
+              isFormulaBroken: false,
+            },
+          },
         },
       });
     });
 
     it('should create a valid formula expression for numeric literals', () => {
       expect(
-        regenerateLayerFromAst(
-          '0',
-          layer,
-          'col1',
-          currentColumn,
+        generateFormulaLayer({
+          formula: '0',
+          layer: {
+            ...layer,
+            columns: {
+              ...layer.columns,
+              col1: currentColumn,
+            },
+          },
+          id: 'col1',
           indexPattern,
-          operationDefinitionMap
-        ).newLayer
+          operations: operationDefinitionMap,
+        })
       ).toEqual({
         ...layer,
         columnOrder: ['col1X0', 'col1'],
@@ -641,14 +655,12 @@ describe('formula', () => {
 
     it('returns the locations of each function', () => {
       expect(
-        regenerateLayerFromAst(
-          'moving_average(average(bytes), window=7) + count()',
+        expandFormulaColumn({
+          formula: 'moving_average(average(bytes), window=7) + count()',
           layer,
-          'col1',
-          currentColumn,
+          id: 'col1',
           indexPattern,
-          operationDefinitionMap
-        ).locations
+        }).meta.locations
       ).toEqual({
         col1X0: { min: 15, max: 29 },
         col1X1: { min: 0, max: 41 },

diff --git a/...ck/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/formula.tsx b/...ck/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/formula.tsx
@@ -11,7 +11,7 @@ import { ReferenceBasedIndexPatternColumn } from '../column_types';
 import { IndexPattern } from '../../../types';
 import { runASTValidation, tryToParse } from './validation';
 import { WrappedFormulaEditor } from './editor';
-import { regenerateLayerFromAst } from './parse';
+import { generateFormulaLayer } from './parse';
 import { generateFormula } from './generate';
 import { filterByVisibleOperation } from './util';
 import { getManagedColumnsFrom } from '../../layer_helpers';
@@ -141,22 +141,20 @@ export const formulaOperation: OperationDefinition<FormulaIndexPatternColumn, 'm
     },
     createCopy(layer, sourceId, targetId, indexPattern, operationDefinitionMap) {
       const currentColumn = layer.columns[sourceId] as FormulaIndexPatternColumn;
-      const tempLayer = {
-        ...layer,
-        columns: {
-          ...layer.columns,
-          [targetId]: { ...currentColumn },
+
+      return generateFormulaLayer({
+        id: targetId,
+        formula: currentColumn.params.formula ?? '',
+        layer: {
+          ...layer,
+          columns: {
+            ...layer.columns,
+            [targetId]: { ...currentColumn },
+          },
         },
-      };
-      const { newLayer } = regenerateLayerFromAst(
-        currentColumn.params.formula ?? '',
-        tempLayer,
-        targetId,
-        currentColumn,
         indexPattern,
-        operationDefinitionMap
-      );
-      return newLayer;
+        operations: operationDefinitionMap,
+      });
     },
 
     paramEditor: WrappedFormulaEditor,

diff --git a/x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/index.ts b/x-pack/plugins/lens/public/indexpattern_datasource/operations/definitions/formula/index.ts
@@ -7,6 +7,8 @@
 
 export type { FormulaIndexPatternColumn } from './formula';
 export { formulaOperation } from './formula';
-export { regenerateLayerFromAst } from './parse';
+
+export { expandFormulaColumn, generateFormulaLayer, generateFormulaColumns } from './parse';
+
 export type { MathIndexPatternColumn } from './math';
 export { mathOperation } from './math';