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

Closes: elastic#103055

x-pack/plugins/lens/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 { regenerateLayerFromFormula, generateFormulaColumns } from '../parse';
 import { filterByVisibleOperation } from '../util';
 import { getColumnTimeShiftWarnings, getDateHistogramInterval } from '../../../../time_shift_utils';
 
@@ -151,16 +151,13 @@ 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) =>
+        regenerateLayerFromFormula({
+          formula: text || '',
+          layer: prevLayer,
           indexPattern,
-          operationDefinitionMap
-        ).newLayer;
-      });
+        })
+      );
     }
   });
 
@@ -173,15 +170,13 @@ 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(
+            regenerateLayerFromFormula({
+              formula: text || '',
+              layer,
+              indexPattern,
+            })
           );
-          updateLayer(newLayer);
         }
 
         return;
@@ -209,14 +204,11 @@ 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 || '',
+              regenerateLayerFromFormula({
+                formula: text || '',
                 layer,
-                columnId,
-                currentColumn,
                 indexPattern,
-                visibleOperationsMap
-              ).newLayer
+              })
             );
           }
         }
@@ -264,15 +256,21 @@ export function FormulaEditor({
         monaco.editor.setModelMarkers(editorModel.current, 'LENS', []);
 
         // Only submit if valid
-        const { newLayer, locations } = regenerateLayerFromAst(
-          text || '',
+        const {
+          columnOrder,
+          columns,
+          meta: { locations },
+        } = generateFormulaColumns({
+          formula: text || '',
           layer,
-          columnId,
-          currentColumn,
           indexPattern,
-          visibleOperationsMap
-        );
-        updateLayer(newLayer);
+        });
+
+        const newLayer = {
+          ...layer,
+          columnOrder,
+          columns,
+        };
 
         const managedColumns = getManagedColumnsFrom(columnId, newLayer.columns);
         const markers: monaco.editor.IMarkerData[] = managedColumns

x-pack/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 { regenerateLayerFromFormula } from './parse';
 import { generateFormula } from './generate';
 import { filterByVisibleOperation } from './util';
 import { getManagedColumnsFrom } from '../../layer_helpers';
@@ -141,22 +141,18 @@ 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 regenerateLayerFromFormula({
+        formula: currentColumn.params.formula ?? '',
+        layer: {
+          ...layer,
+          columns: {
+            ...layer.columns,
+            [targetId]: { ...currentColumn },
+          },
         },
-      };
-      const { newLayer } = regenerateLayerFromAst(
-        currentColumn.params.formula ?? '',
-        tempLayer,
-        targetId,
-        currentColumn,
         indexPattern,
-        operationDefinitionMap
-      );
-      return newLayer;
+      });
     },
 
     paramEditor: WrappedFormulaEditor,

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 { generateFormulaColumns, regenerateLayerFromFormula } from './parse';
+
 export type { MathIndexPatternColumn } from './math';
 export { mathOperation } from './math';

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

@@ -12,6 +12,7 @@ import {
   OperationDefinition,
   GenericOperationDefinition,
   GenericIndexPatternColumn,
+  operationDefinitionMap,
 } from '../index';
 import { IndexPattern, IndexPatternLayer } from '../../../types';
 import { mathOperation } from './math';
@@ -26,7 +27,8 @@ import {
 import { FormulaIndexPatternColumn } from './formula';
 import { getColumnOrder } from '../../layer_helpers';
 
-function getManagedId(mainId: string, index: number) {
+/** @internal **/
+export function getManagedId(mainId: string, index: number) {
   return `${mainId}X${index}`;
 }
 
@@ -35,15 +37,15 @@ function parseAndExtract(
   layer: IndexPatternLayer,
   columnId: string,
   indexPattern: IndexPattern,
-  operationDefinitionMap: Record<string, GenericOperationDefinition>,
+  operations: Record<string, GenericOperationDefinition>,
   label?: string
 ) {
-  const { root, error } = tryToParse(text, operationDefinitionMap);
+  const { root, error } = tryToParse(text, operations);
   if (error || root == null) {
     return { extracted: [], isValid: false };
   }
   // before extracting the data run the validation task and throw if invalid
-  const errors = runASTValidation(root, layer, indexPattern, operationDefinitionMap);
+  const errors = runASTValidation(root, layer, indexPattern, operations);
   if (errors.length) {
     return { extracted: [], isValid: false };
   }
@@ -52,7 +54,7 @@ function parseAndExtract(
     */
   const extracted = extractColumns(
     columnId,
-    operationDefinitionMap,
+    operations,
     root,
     layer,
     indexPattern,
@@ -185,63 +187,88 @@ function extractColumns(
   return columns;
 }
 
-export function regenerateLayerFromAst(
-  text: string,
-  layer: IndexPatternLayer,
-  columnId: string,
-  currentColumn: FormulaIndexPatternColumn,
-  indexPattern: IndexPattern,
-  operationDefinitionMap: Record<string, GenericOperationDefinition>
-) {
-  const { extracted, isValid } = parseAndExtract(
-    text,
-    layer,
-    columnId,
-    indexPattern,
-    filterByVisibleOperation(operationDefinitionMap),
-    currentColumn.customLabel ? currentColumn.label : undefined
-  );
+/** @public **/
+export function generateFormulaColumns({
+  formula,
+  layer,
+  indexPattern,
+}: {
+  formula: string;
+  layer: IndexPatternLayer;
+  indexPattern: IndexPattern;
+}) {
+  const columnEntries = Object.entries(layer.columns);
+  const formulaColumnsIds = columnEntries
+    .filter(([, value]) => value.operationType === 'formula')
+    .map(([key]) => key);
 
-  const columns = { ...layer.columns };
+  const { columns, meta } = Object.entries(layer.columns).reduce<{
+    columns: Record<string, GenericIndexPatternColumn>;
+    meta: {
+      locations: Record<string, TinymathLocation>;
+    };
+  }>(
+    (acc, [id, value]) => {
+      if (formulaColumnsIds.includes(id)) {
+        const { extracted, isValid } = parseAndExtract(
+          formula,
+          layer,
+          id,
+          indexPattern,
+          filterByVisibleOperation(operationDefinitionMap),
+          value.label
+        );
 
-  const locations: Record<string, TinymathLocation> = {};
+        extracted.forEach(({ column, location }, index) => {
+          const managedId = getManagedId(id, index);
+          acc.columns[managedId] = column;
 
-  Object.keys(columns).forEach((k) => {
-    if (k.startsWith(columnId)) {
-      delete columns[k];
-    }
-  });
-
-  extracted.forEach(({ column, location }, index) => {
-    columns[getManagedId(columnId, index)] = column;
-    if (location) locations[getManagedId(columnId, index)] = location;
-  });
-
-  columns[columnId] = {
-    ...currentColumn,
-    label: !currentColumn.customLabel
-      ? text ??
-        i18n.translate('xpack.lens.indexPattern.formulaLabel', {
-          defaultMessage: 'Formula',
-        })
-      : currentColumn.label,
-    params: {
-      ...currentColumn.params,
-      formula: text,
-      isFormulaBroken: !isValid,
+          if (location) {
+            acc.meta.locations[managedId] = location;
+          }
+        });
+
+        acc.columns[id] = {
+          ...value,
+          label:
+            value.label ??
+            (formula ||
+              i18n.translate('xpack.lens.indexPattern.formulaLabel', {
+                defaultMessage: 'Formula',
+              })),
+          references: !isValid ? [] : [getManagedId(id, extracted.length - 1)],
+          params: {
+            // ? @todo should be checked
+            // ...value.params,
+            formula,
+            isFormulaBroken: !isValid,
+          },
+        } as FormulaIndexPatternColumn;
+      } else {
+        acc.columns[id] = { ...value };
+      }
+      return acc;
     },
-    references: !isValid ? [] : [getManagedId(columnId, extracted.length - 1)],
-  } as FormulaIndexPatternColumn;
+    {
+      columns: {},
+      meta: {
+        locations: {},
+      },
+    }
+  );
 
   return {
-    newLayer: {
+    columns,
+    columnOrder: getColumnOrder({
       ...layer,
       columns,
-      columnOrder: getColumnOrder({
-        ...layer,
-        columns,
-      }),
-    },
-    locations,
+    }),
+    meta,
   };
 }
+
+/** @internal **/
+export function regenerateLayerFromFormula(params: Parameters<typeof generateFormulaColumns>[0]) {
+  const { columns, columnOrder } = generateFormulaColumns(params);
+  return { ...params.layer, columns, columnOrder };
+}

x-pack/plugins/lens/public/indexpattern_datasource/operations/layer_helpers.ts

@@ -36,7 +36,7 @@ import {
   ReferenceBasedIndexPatternColumn,
   BaseIndexPatternColumn,
 } from './definitions/column_types';
-import { FormulaIndexPatternColumn, regenerateLayerFromAst } from './definitions/formula';
+import { FormulaIndexPatternColumn, regenerateLayerFromFormula } from './definitions/formula';
 import type { TimeScaleUnit } from '../../../common/expressions';
 import { isColumnOfType } from './definitions/helpers';
 
@@ -533,14 +533,11 @@ export function replaceColumn({
 
       try {
         newLayer = newColumn.params.formula
-          ? regenerateLayerFromAst(
-              newColumn.params.formula,
-              basicLayer,
-              columnId,
-              newColumn,
+          ? regenerateLayerFromFormula({
+              formula: newColumn.params.formula,
+              layer: basicLayer,
               indexPattern,
-              operationDefinitionMap
-            ).newLayer
+            })
           : basicLayer;
       } catch (e) {
         newLayer = basicLayer;