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

Closes: elastic#103055

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 { 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,10 @@ import type {
  FieldFormatsSetup,
 } from '../../../../../src/plugins/field_formats/public';
 
+export { generateFormulaColumns };
+
+export type GenerateFormulaColumnFn = typeof generateFormulaColumns;
+
 export interface IndexPatternDatasourceSetupPlugins {
  expressions: ExpressionsSetup;
  fieldFormats: FieldFormatsSetup;

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,15 @@ 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({
+ id: columnId,
+ label: currentColumn.label,
+ formula: text || '',
+ layer: prevLayer,
  indexPattern,
- operationDefinitionMap
- ).newLayer;
- });
+ })
+ );
  }
  });
 
@@ -173,15 +172,15 @@ 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({
+ id: columnId,
+ label: currentColumn.label,
+ formula: text || '',
+ layer,
+ indexPattern,
+ })
  );
- updateLayer(newLayer);
  }
 
  return;
@@ -209,14 +208,13 @@ 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({
+ id: columnId,
+ label: currentColumn.label,
+ formula: text || '',
  layer,
- columnId,
- currentColumn,
  indexPattern,
- visibleOperationsMap
- ).newLayer
+ })
  );
  }
  }
@@ -264,15 +262,23 @@ export function FormulaEditor({
  monaco.editor.setModelMarkers(editorModel.current, 'LENS', []);
 
  // Only submit if valid
- const { newLayer, locations } = regenerateLayerFromAst(
- text || '',
+ const {
+ columnOrder,
+ columns,
+ meta: { locations },
+ } = generateFormulaColumns({
+ id: columnId,
+ label: currentColumn.label,
+ 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,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 regenerateLayerFromFormula({
+ id: targetId,
+ formula: currentColumn.params.formula ?? '',
+ layer: {
+ ...layer,
+ columns: {
+ ...layer.columns,
+ [targetId]: { ...currentColumn },
+ },
  },
- };
- const { newLayer } = regenerateLayerFromAst(
- currentColumn.params.formula ?? '',
- tempLayer,
- targetId,
- currentColumn,
+ label: currentColumn.label,
  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,86 @@ 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,
+ label,
+ id,
+}: {
+ id: string;
+ formula: string;
+ layer: IndexPatternLayer;
+ indexPattern: IndexPattern;
+ label?: string;
+}) {
+ const { columns, meta } = Object.entries(layer.columns).reduce<{
+ columns: Record<string, GenericIndexPatternColumn>;
+ meta: {
+ locations: Record<string, TinymathLocation>;
+ };
+ }>(
+ (acc, [currentColumnId, currentColumn]) => {
+ if (currentColumnId === id) {
+ const { extracted, isValid } = parseAndExtract(
+ formula,
+ layer,
+ currentColumnId,
+ indexPattern,
+ filterByVisibleOperation(operationDefinitionMap),
+ label
+ );
 
- const columns = { ...layer.columns };
+ extracted.forEach(({ column, location }, index) => {
+ const managedId = getManagedId(currentColumnId, index);
+ acc.columns[managedId] = column;
 
- const locations: Record<string, TinymathLocation> = {};
+ if (location) {
+ acc.meta.locations[managedId] = location;
+ }
+ });
 
- 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,
+ acc.columns[currentColumnId] = {
+ ...currentColumn,
+ label:
+ currentColumn.label ??
+ (formula ||
+ i18n.translate('xpack.lens.indexPattern.formulaLabel', {
+ defaultMessage: 'Formula',
+ })),
+ references: !isValid ? [] : [getManagedId(currentColumnId, extracted.length - 1)],
+ params: {
+ // ? @todo should be checked
+ // ...value.params,
+ formula,
+ isFormulaBroken: !isValid,
+ },
+ } as FormulaIndexPatternColumn;
+ } else {
+ acc.columns[currentColumnId] = { ...currentColumn };
+ }
+ 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/index.ts

@@ -8,6 +8,7 @@
 export * from './operations';
 export * from './layer_helpers';
 export * from './time_scale_utils';
+
 export type {
  OperationType,
  BaseIndexPatternColumn,