[Obs ai assistant][ESQL] Visualizes a query

packages/kbn-visualization-utils/index.ts

@@ -7,3 +7,4 @@
  */
 
 export { getTimeZone } from './src/get_timezone';
+export { getLensAttributes } from './src/get_lens_attributes';

packages/kbn-visualization-utils/src/get_lens_attributes.ts

@@ -0,0 +1,99 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+import { i18n } from '@kbn/i18n';
+import type { Ast } from '@kbn/interpreter';
+import type { IconType } from '@elastic/eui/src/components/icon/icon';
+import type { DataView } from '@kbn/data-views-plugin/public';
+import type { AggregateQuery, Query, Filter } from '@kbn/es-query';
+
+/**
+ * Indicates what was changed in this table compared to the currently active table of this layer.
+ * * `initial` means the layer associated with this table does not exist in the current configuration
+ * * `unchanged` means the table is the same in the currently active configuration
+ * * `reduced` means the table is a reduced version of the currently active table (some columns dropped, but not all of them)
+ * * `extended` means the table is an extended version of the currently active table (added one or multiple additional columns)
+ * * `reorder` means the table columns have changed order, which change the data as well
+ * * `layers` means the change is a change to the layer structure, not to the table
+ */
+export type TableChangeType =
+  | 'initial'
+  | 'unchanged'
+  | 'reduced'
+  | 'extended'
+  | 'reorder'
+  | 'layers';
+
+export interface Suggestion<T = unknown, V = unknown> {
+  visualizationId: string;
+  datasourceState?: V;
+  datasourceId?: string;
+  columns: number;
+  score: number;
+  title: string;
+  visualizationState: T;
+  previewExpression?: Ast | string;
+  previewIcon: IconType;
+  hide?: boolean;
+  // flag to indicate if the visualization is incomplete
+  incomplete?: boolean;
+  changeType: TableChangeType;
+  keptLayerIds: string[];
+}
+
+export const getLensAttributes = ({
+  filters,
+  query,
+  suggestion,
+  dataView,
+}: {
+  filters: Filter[];
+  query: Query | AggregateQuery;
+  suggestion: Suggestion | undefined;
+  dataView?: DataView;
+}) => {
+  const suggestionDatasourceState = Object.assign({}, suggestion?.datasourceState);
+  const suggestionVisualizationState = Object.assign({}, suggestion?.visualizationState);
+  const datasourceStates =
+    suggestion && suggestion.datasourceState
+      ? {
+          [suggestion.datasourceId!]: {
+            ...suggestionDatasourceState,
+          },
+        }
+      : {
+          formBased: {},
+        };
+  const visualization = suggestionVisualizationState;
+  const attributes = {
+    title: suggestion
+      ? suggestion.title
+      : i18n.translate('xpack.lens.config.suggestion.title', {
+          defaultMessage: 'New suggestion',
+        }),
+    references: [
+      {
+        id: dataView?.id ?? '',
+        name: `textBasedLanguages-datasource-layer-suggestion`,
+        type: 'index-pattern',
+      },
+    ],
+    state: {
+      datasourceStates,
+      filters,
+      query,
+      visualization,
+      ...(dataView &&
+        dataView.id &&
+        !dataView.isPersisted() && {
+          adHocDataViews: { [dataView.id]: dataView.toSpec(false) },
+        }),
+    },
+    visualizationType: suggestion ? suggestion.visualizationId : 'lnsXY',
+  };
+  return attributes;
+};

x-pack/plugins/lens/public/async_services.ts

@@ -51,3 +51,4 @@ export * from './chart_info_api';
 
 export * from './trigger_actions/open_in_discover_helpers';
 export * from './trigger_actions/open_lens_config/helpers';
+export * from './trigger_actions/open_lens_config/embeddable_edit_action_helpers';

x-pack/plugins/lens/public/embeddable/embeddable.tsx

@@ -750,7 +750,8 @@ export class Embeddable
   async updateVisualization(
     datasourceState: unknown,
     visualizationState: unknown,
-    visualizationType?: string
+    visualizationType?: string,
+    onUpdate?: (input: TypedLensByValueInput['attributes']) => void
   ) {
     const viz = this.savedVis;
     const activeDatasourceId = (this.activeDatasourceId ??
@@ -785,13 +786,14 @@ export class Embeddable
         },
         references,
         visualizationType: visualizationType ?? viz.visualizationType,
-      };
+      } as TypedLensByValueInput['attributes'];
 
       /**
        * SavedObjectId is undefined for by value panels and defined for the by reference ones.
        * Here we are converting the by reference panels to by value when user is inline editing
        */
       this.updateInput({ attributes: attrs, savedObjectId: undefined });
+      onUpdate?.(attrs);
       /**
        * Should load again the user messages,
        * otherwise the embeddable state is stuck in an error state
@@ -800,13 +802,17 @@ export class Embeddable
     }
   }
 
-  async updateSuggestion(attrs: LensSavedObjectAttributes) {
+  async updateSuggestion(
+    attrs: LensSavedObjectAttributes,
+    onUpdate?: (input: TypedLensByValueInput['attributes']) => void
+  ) {
     const viz = this.savedVis;
     const newViz = {
       ...viz,
       ...attrs,
-    };
+    } as TypedLensByValueInput['attributes'];
     this.updateInput({ attributes: newViz });
+    onUpdate?.(newViz);
   }
 
   /**
@@ -844,7 +850,10 @@ export class Embeddable
     this.updateInput({ attributes: attrs, savedObjectId });
   }
 
-  async openConfingPanel(startDependencies: LensPluginStartDependencies) {
+  async openConfingPanel(
+    startDependencies: LensPluginStartDependencies,
+    onUpdate?: (input: TypedLensByValueInput['attributes']) => void
+  ) {
     const { getEditLensConfiguration } = await import('../async_services');
     const Component = await getEditLensConfiguration(
       this.deps.coreStart,
@@ -861,8 +870,10 @@ export class Embeddable
       return (
         <Component
           attributes={attributes}
-          updatePanelState={this.updateVisualization.bind(this)}
-          updateSuggestion={this.updateSuggestion.bind(this)}
+          updatePanelState={(...attrs) => {
+            this.updateVisualization.bind(this)(...attrs, onUpdate);
+          }}
+          updateSuggestion={(...attrs) => this.updateSuggestion.bind(this)(...attrs, onUpdate)}
           datasourceId={datasourceId}
           lensAdapters={this.lensInspector.adapters}
           output$={this.getOutput$()}

x-pack/plugins/lens/public/plugin.ts

@@ -105,11 +105,13 @@ import type {
 } from './types';
 import { getLensAliasConfig } from './vis_type_alias';
 import { createOpenInDiscoverAction } from './trigger_actions/open_in_discover_action';
+import { EditLensEmbeddableAction } from './trigger_actions/open_lens_config/embeddable_edit_action';
 import { ConfigureInLensPanelAction } from './trigger_actions/open_lens_config/action';
 import { visualizeFieldAction } from './trigger_actions/visualize_field_actions';
 import { visualizeTSVBAction } from './trigger_actions/visualize_tsvb_actions';
 import { visualizeAggBasedVisAction } from './trigger_actions/visualize_agg_based_vis_actions';
 import { visualizeDashboardVisualizePanelction } from './trigger_actions/dashboard_visualize_panel_actions';
+import { inAppEditTrigger } from './trigger_actions/open_lens_config/in_app_edit_trigger';
 
 import type { LensEmbeddableInput } from './embeddable';
 import { EmbeddableFactory, LensEmbeddableStartServices } from './embeddable/embeddable_factory';
@@ -327,6 +329,10 @@ export class LensPlugin {
         this.editorFrameService!.loadVisualizations(),
         this.editorFrameService!.loadDatasources(),
       ]);
+      const { setVisualizationMap, setDatasourceMap } = await import('./async_services');
+      setDatasourceMap(datasourceMap);
+      setVisualizationMap(visualizationMap);
+
       const eventAnnotationService = await plugins.eventAnnotation.getService();
 
       if (plugins.usageCollection) {
@@ -598,6 +604,13 @@ export class LensPlugin {
     );
     startDependencies.uiActions.addTriggerAction('CONTEXT_MENU_TRIGGER', editInLensAction);
 
+    // this trigger enables external consumers to use the inline editing flyout
+    startDependencies.uiActions.registerTrigger(inAppEditTrigger);
+
+    // Allows the Lens embeddable to easily open the inapp editing flyout
+    const editLensEmbeddableAction = new EditLensEmbeddableAction(startDependencies, core);
+    startDependencies.uiActions.addTriggerAction('IN_APP_EDIT_TRIGGER', editLensEmbeddableAction);
+
     const discoverLocator = startDependencies.share?.url.locators.get('DISCOVER_APP_LOCATOR');
     if (discoverLocator) {
       startDependencies.uiActions.addTriggerAction(

x-pack/plugins/lens/public/trigger_actions/open_lens_config/embeddable_edit_action.tsx

@@ -0,0 +1,61 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { i18n } from '@kbn/i18n';
+import type { CoreStart } from '@kbn/core/public';
+import { Action } from '@kbn/ui-actions-plugin/public';
+import type { LensPluginStartDependencies } from '../../plugin';
+import type { TypedLensByValueInput } from '../../embeddable/embeddable_component';
+
+const ACTION_EDIT_LENS_EMBEDDABLE = 'ACTION_EDIT_LENS_EMBEDDABLE';
+
+interface Context {
+  attributes?: TypedLensByValueInput['attributes'];
+  id?: string;
+  onUpdate?: (input: TypedLensByValueInput['attributes']) => void;
+}
+
+export const getAsyncHelpers = async () => await import('../../async_services');
+
+export class EditLensEmbeddableAction implements Action<Context> {
+  public type = ACTION_EDIT_LENS_EMBEDDABLE;
+  public id = ACTION_EDIT_LENS_EMBEDDABLE;
+  public order = 50;
+
+  constructor(
+    protected readonly startDependencies: LensPluginStartDependencies,
+    protected readonly core: CoreStart
+  ) {}
+
+  public getDisplayName(): string {
+    return i18n.translate('xpack.lens.app.editLensEmbeddableLabel', {
+      defaultMessage: 'Edit visualization',
+    });
+  }
+
+  public getIconType() {
+    return 'pencil';
+  }
+
+  public async isCompatible() {
+    // compatible only when ES|QL advanced setting is enabled
+    const { isEmbeddableEditActionCompatible } = await getAsyncHelpers();
+    return isEmbeddableEditActionCompatible(this.core);
+  }
+
+  public async execute({ attributes, id, onUpdate }: Context) {
+    const { executeEditEmbeddableAction } = await getAsyncHelpers();
+    if (attributes) {
+      executeEditEmbeddableAction({
+        deps: this.startDependencies,
+        core: this.core,
+        attributes,
+        embeddableId: id,
+        onUpdate,
+      });
+    }
+  }
+}

x-pack/plugins/lens/public/trigger_actions/open_lens_config/embeddable_edit_action_helpers.ts

@@ -0,0 +1,70 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { createGetterSetter } from '@kbn/kibana-utils-plugin/common';
+import type { CoreStart } from '@kbn/core/public';
+import { IncompatibleActionError } from '@kbn/ui-actions-plugin/public';
+import type { Datasource, Visualization } from '../../types';
+import type { LensPluginStartDependencies } from '../../plugin';
+import { generateId } from '../../id_generator';
+import { executeAction } from './helpers';
+import type { TypedLensByValueInput } from '../../embeddable/embeddable_component';
+
+// datasourceMap and visualizationMap setters/getters
+export const [getVisualizationMap, setVisualizationMap] = createGetterSetter<
+  Record<string, Visualization<unknown, unknown, unknown>>
+>('VisualizationMap', false);
+
+export const [getDatasourceMap, setDatasourceMap] = createGetterSetter<
+  Record<string, Datasource<unknown, unknown>>
+>('DatasourceMap', false);
+
+export function isEmbeddableEditActionCompatible(core: CoreStart) {
+  return core.uiSettings.get('discover:enableESQL');
+}
+
+export async function executeEditEmbeddableAction({
+  deps,
+  core,
+  attributes,
+  embeddableId,
+  onUpdate,
+}: {
+  deps: LensPluginStartDependencies;
+  core: CoreStart;
+  attributes: TypedLensByValueInput['attributes'];
+  embeddableId?: string;
+  onUpdate?: (input: TypedLensByValueInput['attributes']) => void;
+}) {
+  const isCompatibleAction = isEmbeddableEditActionCompatible(core);
+  const defaultDataView = await deps.dataViews.getDefaultDataView({
+    displayErrors: false,
+  });
+  if (!isCompatibleAction || !defaultDataView) {
+    throw new IncompatibleActionError();
+  }
+
+  const input = {
+    attributes,
+    id: embeddableId ?? generateId(),
+  };
+  const embeddableStart = deps.embeddable;
+  const factory = embeddableStart.getEmbeddableFactory('lens');
+  if (!factory) {
+    return undefined;
+  }
+  const embeddable = await factory.create(input);
+  // open the flyout if embeddable has been created successfully
+  if (embeddable) {
+    executeAction({
+      embeddable,
+      startDependencies: deps,
+      overlays: core.overlays,
+      theme: core.theme,
+      onUpdate,
+    });
+  }
+}

x-pack/plugins/lens/public/trigger_actions/open_lens_config/helpers.ts

@@ -12,12 +12,14 @@ import { toMountPoint } from '@kbn/kibana-react-plugin/public';
 import { IncompatibleActionError } from '@kbn/ui-actions-plugin/public';
 import { isLensEmbeddable } from '../utils';
 import type { LensPluginStartDependencies } from '../../plugin';
+import type { TypedLensByValueInput } from '../../embeddable/embeddable_component';
 
 interface Context {
   embeddable: IEmbeddable;
   startDependencies: LensPluginStartDependencies;
   overlays: OverlayStart;
   theme: ThemeServiceStart;
+  onUpdate?: (input: TypedLensByValueInput['attributes']) => void;
 }
 
 export async function isActionCompatible(embeddable: IEmbeddable) {
@@ -26,14 +28,20 @@ export async function isActionCompatible(embeddable: IEmbeddable) {
   return Boolean(isLensEmbeddable(embeddable) && embeddable.getIsEditable() && inDashboardEditMode);
 }
 
-export async function executeAction({ embeddable, startDependencies, overlays, theme }: Context) {
+export async function executeAction({
+  embeddable,
+  startDependencies,
+  overlays,
+  theme,
+  onUpdate,
+}: Context) {
   const isCompatibleAction = await isActionCompatible(embeddable);
   if (!isCompatibleAction || !isLensEmbeddable(embeddable)) {
     throw new IncompatibleActionError();
   }
   const rootEmbeddable = embeddable.getRoot();
   const overlayTracker = tracksOverlays(rootEmbeddable) ? rootEmbeddable : undefined;
-  const ConfigPanel = await embeddable.openConfingPanel(startDependencies);
+  const ConfigPanel = await embeddable.openConfingPanel(startDependencies, onUpdate);
   if (ConfigPanel) {
     const handle = overlays.openFlyout(
       toMountPoint(