[Lens] removes DatasourcePublicAPI (#170725)

## Summary We have 2 very similar structures in Lens - FramePublicAPI and DatasourcePublicAPI. Creating `DatasourcePublicAPI` was an attempt to remove framePublicAPI by firstly limiting it and then just using redux store directly. Unfortunately, due to Lens architecture it never progressed so we're stuck with those two structures for no reason whatsover. This PR merges them back together. Apart from that it also moves some code from app.tsx to the hook that is responsible for creating `getUserMessages` helper. This way we could use it in other places (inline editor, for example) ### Checklist Delete any items that are not applicable to this PR. - [ ] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md) - [ ] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [ ] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios - [ ] Any UI touched in this PR is usable by keyboard only (learn more about [keyboard accessibility](https://webaim.org/techniques/keyboard/)) - [ ] Any UI touched in this PR does not create any new axe failures (run axe in browser: [FF](https://addons.mozilla.org/en-US/firefox/addon/axe-devtools/), [Chrome](https://chrome.google.com/webstore/detail/axe-web-accessibility-tes/lhdoppojpmngadmnindnejefpokejbdd?hl=en-US)) - [ ] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [ ] This renders correctly on smaller devices using a responsive layout. (You can test this [in your browser](https://www.browserstack.com/guide/responsive-testing-on-local-server)) - [ ] This was checked for [cross-browser compatibility](https://www.elastic.co/support/matrix#matrix_browsers) ### Risk Matrix Delete this section if it is not applicable to this PR. Before closing this PR, invite QA, stakeholders, and other developers to identify risks that should be tested prior to the change/feature release. When forming the risk matrix, consider some of the following examples and how they may potentially impact the change: | Risk | Probability | Severity | Mitigation/Notes | |---------------------------|-------------|----------|-------------------------| | Multiple Spaces—unexpected behavior in non-default Kibana Space. | Low | High | Integration tests will verify that all features are still supported in non-default Kibana Space and when user switches between spaces. | | Multiple nodes—Elasticsearch polling might have race conditions when multiple Kibana nodes are polling for the same tasks. | High | Low | Tasks are idempotent, so executing them multiple times will not result in logical error, but will degrade performance. To test for this case we add plenty of unit tests around this logic and document manual testing procedure. | | Code should gracefully handle cases when feature X or plugin Y are disabled. | Medium | High | Unit tests will verify that any feature flag or plugin combination still results in our service operational. | | [See more potential risk examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx) | ### For maintainers - [ ] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)
2025-04-23 09:19:04 -04:00 · 2023-11-09 11:09:52 +01:00 · 2023-11-09 11:09:52 +01:00 · 5fa46caa5c
commit 5fa46caa5c
parent 0b40e3e7c1
26 changed files with 289 additions and 273 deletions
--- a/x-pack/plugins/lens/public/app_plugin/app.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.tsx
@ -17,7 +17,7 @@ import type { LensAppLocatorParams } from '../../common/locator/locator';
 import { LensAppProps, LensAppServices } from './types';
 import { LensTopNavMenu } from './lens_top_nav';
 import { LensByReferenceInput } from '../embeddable';
-import { AddUserMessages, EditorFrameInstance, UserMessage, UserMessagesGetter } from '../types';
+import { AddUserMessages, EditorFrameInstance, UserMessagesGetter } from '../types';
 import { Document } from '../persistence/saved_object_store';

 import {
@ -28,9 +28,8 @@ import {
  LensAppState,
  selectSavedObjectFormat,
  updateIndexPatterns,
-  updateDatasourceState,
  selectActiveDatasourceId,
-  selectFrameDatasourceAPI,
+  selectFramePublicAPI,
 } from '../state_management';
 import { SaveModalContainer, runSaveLensVisualization } from './save_modal_container';
 import { LensInspector } from '../lens_inspector_service';
@ -41,10 +40,7 @@ import {
  createIndexPatternService,
 } from '../data_views_service/service';
 import { replaceIndexpattern } from '../state_management/lens_slice';
-import {
-  filterAndSortUserMessages,
-  getApplicationUserMessages,
-} from './get_application_user_messages';
+import { useApplicationUserMessages } from './get_application_user_messages';

 export type SaveProps = Omit<OnSaveProps, 'onTitleDuplicate' | 'newDescription'> & {
  returnToOrigin: boolean;
@ -509,99 +505,25 @@ export function App({

  const activeDatasourceId = useLensSelector(selectActiveDatasourceId);

-  const frameDatasourceAPI = useLensSelector((state) =>
-    selectFrameDatasourceAPI(state, datasourceMap)
-  );
+  const framePublicAPI = useLensSelector((state) => selectFramePublicAPI(state, datasourceMap));

-  const [userMessages, setUserMessages] = useState<UserMessage[]>([]);
-
-  useEffect(() => {
-    setUserMessages([
-      ...(activeDatasourceId
-        ? datasourceMap[activeDatasourceId].getUserMessages(
-            datasourceStates[activeDatasourceId].state,
-            {
-              frame: frameDatasourceAPI,
-              setState: (newStateOrUpdater) => {
-                dispatch(
-                  updateDatasourceState({
-                    newDatasourceState:
-                      typeof newStateOrUpdater === 'function'
-                        ? newStateOrUpdater(datasourceStates[activeDatasourceId].state)
-                        : newStateOrUpdater,
-                    datasourceId: activeDatasourceId,
-                  })
-                );
-              },
-            }
-          )
-        : []),
-      ...(visualization.activeId && visualization.state
-        ? visualizationMap[visualization.activeId]?.getUserMessages?.(visualization.state, {
-            frame: frameDatasourceAPI,
-          }) ?? []
-        : []),
-      ...getApplicationUserMessages({
-        visualizationType: persistedDoc?.visualizationType,
-        visualizationMap,
-        visualization,
-        activeDatasource: activeDatasourceId ? datasourceMap[activeDatasourceId] : null,
-        activeDatasourceState: activeDatasourceId ? datasourceStates[activeDatasourceId] : null,
-        core: coreStart,
-        dataViews: frameDatasourceAPI.dataViews,
-      }),
-    ]);
-  }, [
-    activeDatasourceId,
+  const { getUserMessages, addUserMessages } = useApplicationUserMessages({
    coreStart,
-    datasourceMap,
-    datasourceStates,
+    framePublicAPI,
+    activeDatasourceId,
+    datasourceState:
+      activeDatasourceId && datasourceStates[activeDatasourceId]
+        ? datasourceStates[activeDatasourceId]
+        : null,
+    datasource:
+      activeDatasourceId && datasourceMap[activeDatasourceId]
+        ? datasourceMap[activeDatasourceId]
+        : null,
    dispatch,
-    frameDatasourceAPI,
-    persistedDoc?.visualizationType,
-    visualization,
-    visualizationMap,
-  ]);
-
-  // these are messages managed from other parts of Lens
-  const [additionalUserMessages, setAdditionalUserMessages] = useState<Record<string, UserMessage>>(
-    {}
-  );
-
-  const getUserMessages: UserMessagesGetter = (locationId, filterArgs) =>
-    filterAndSortUserMessages(
-      [...userMessages, ...Object.values(additionalUserMessages)],
-      locationId,
-      filterArgs ?? {}
-    );
-
-  const addUserMessages: AddUserMessages = (messages) => {
-    const newMessageMap = {
-      ...additionalUserMessages,
-    };
-
-    const addedMessageIds: string[] = [];
-    messages.forEach((message) => {
-      if (!newMessageMap[message.uniqueId]) {
-        addedMessageIds.push(message.uniqueId);
-        newMessageMap[message.uniqueId] = message;
-      }
-    });
-
-    if (addedMessageIds.length) {
-      setAdditionalUserMessages(newMessageMap);
-    }
-
-    return () => {
-      const withMessagesRemoved = {
-        ...additionalUserMessages,
-      };
-
-      addedMessageIds.forEach((id) => delete withMessagesRemoved[id]);
-
-      setAdditionalUserMessages(withMessagesRemoved);
-    };
-  };
+    visualization: visualization.activeId ? visualizationMap[visualization.activeId] : undefined,
+    visualizationType: visualization.activeId,
+    visualizationState: visualization,
+  });

  return (
    <>
--- a/x-pack/plugins/lens/public/app_plugin/get_application_user_messages.test.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/get_application_user_messages.test.tsx
@ -25,8 +25,8 @@ describe('application-level user messages', () => {
      getApplicationUserMessages({
        visualizationType: undefined,

-        visualizationMap: {},
-        visualization: { activeId: '', state: {} },
+        visualization: undefined,
+        visualizationState: { activeId: '', state: {} },
        activeDatasource: {} as Datasource,
        activeDatasourceState: null,
        dataViews: {} as DataViewsState,
@ -53,8 +53,8 @@ describe('application-level user messages', () => {
    expect(
      getApplicationUserMessages({
        visualizationType: '123',
-        visualizationMap: {},
-        visualization: { activeId: 'id_for_type_that_doesnt_exist', state: {} },
+        visualization: undefined,
+        visualizationState: { activeId: 'id_for_type_that_doesnt_exist', state: {} },

        activeDatasource: {} as Datasource,
        activeDatasourceState: null,
@ -84,8 +84,8 @@ describe('application-level user messages', () => {
        activeDatasource: null,

        visualizationType: '123',
-        visualizationMap: { 'some-id': {} as Visualization },
-        visualization: { activeId: 'some-id', state: {} },
+        visualization: {} as Visualization,
+        visualizationState: { activeId: 'some-id', state: {} },
        activeDatasourceState: null,
        dataViews: {} as DataViewsState,
        core: {} as CoreStart,
@ -138,8 +138,8 @@ describe('application-level user messages', () => {

    const irrelevantProps = {
      dataViews: {} as DataViewsState,
-      visualizationMap: { foo: {} as Visualization },
-      visualization: { activeId: 'foo', state: {} },
+      visualization: {} as Visualization,
+      visualizationState: { activeId: 'foo', state: {} },
    };

    it('generates error if missing an index pattern', () => {
--- a/x-pack/plugins/lens/public/app_plugin/get_application_user_messages.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/get_application_user_messages.tsx
@ -5,18 +5,27 @@
 * 2.0.
 */

-import React from 'react';
+import React, { useEffect, useState } from 'react';
 import { i18n } from '@kbn/i18n';
 import { RedirectAppLinks } from '@kbn/shared-ux-link-redirect-app';
 import { FormattedMessage } from '@kbn/i18n-react';
 import type { CoreStart } from '@kbn/core/public';
-import type { DataViewsState, VisualizationState } from '../state_management';
+import { Dispatch } from '@reduxjs/toolkit';
+import {
+  updateDatasourceState,
+  type DataViewsState,
+  type VisualizationState,
+  DatasourceState,
+} from '../state_management';
 import type {
+  AddUserMessages,
  Datasource,
+  FramePublicAPI,
  UserMessage,
  UserMessageFilters,
  UserMessagesDisplayLocationId,
-  VisualizationMap,
+  UserMessagesGetter,
+  Visualization,
 } from '../types';
 import { getMissingIndexPattern } from '../editor_frame_service/editor_frame/state_helpers';

@ -26,15 +35,15 @@ import { getMissingIndexPattern } from '../editor_frame_service/editor_frame/sta
 export const getApplicationUserMessages = ({
  visualizationType,
  visualization,
-  visualizationMap,
+  visualizationState,
  activeDatasource,
  activeDatasourceState,
  dataViews,
  core,
 }: {
  visualizationType: string | null | undefined;
-  visualization: VisualizationState | undefined;
-  visualizationMap: VisualizationMap;
+  visualization: Visualization | undefined;
+  visualizationState: VisualizationState | undefined;
  activeDatasource: Datasource | null | undefined;
  activeDatasourceState: { isLoading: boolean; state: unknown } | null;
  dataViews: DataViewsState;
@ -46,8 +55,8 @@ export const getApplicationUserMessages = ({
    messages.push(getMissingVisTypeError());
  }

-  if (visualization?.activeId && !visualizationMap[visualization.activeId]) {
-    messages.push(getUnknownVisualizationTypeError(visualization.activeId));
+  if (visualizationState?.activeId && !visualization) {
+    messages.push(getUnknownVisualizationTypeError(visualizationState.activeId));
  }

  if (!activeDatasource) {
@ -182,8 +191,8 @@ function getMissingIndexPatternsErrors(

 export const filterAndSortUserMessages = (
  userMessages: UserMessage[],
-  locationId: UserMessagesDisplayLocationId | UserMessagesDisplayLocationId[] | undefined,
-  { dimensionId, severity }: UserMessageFilters
+  locationId?: UserMessagesDisplayLocationId | UserMessagesDisplayLocationId[],
+  { dimensionId, severity }: UserMessageFilters = {}
 ) => {
  const locationIds = Array.isArray(locationId)
    ? locationId
@ -231,3 +240,112 @@ function bySeverity(a: UserMessage, b: UserMessage) {
  }
  return 1;
 }
+
+export const useApplicationUserMessages = ({
+  coreStart,
+  dispatch,
+  activeDatasourceId,
+  datasource,
+  datasourceState,
+  framePublicAPI,
+  visualizationType,
+  visualization,
+  visualizationState,
+}: {
+  activeDatasourceId: string | null;
+  coreStart: CoreStart;
+  datasource: Datasource | null;
+  datasourceState: DatasourceState | null;
+  dispatch: Dispatch;
+  framePublicAPI: FramePublicAPI;
+  visualizationType: string | null;
+  visualizationState?: VisualizationState;
+  visualization?: Visualization;
+}) => {
+  const [userMessages, setUserMessages] = useState<UserMessage[]>([]);
+  // these are messages managed from other parts of Lens
+  const [additionalUserMessages, setAdditionalUserMessages] = useState<Record<string, UserMessage>>(
+    {}
+  );
+
+  useEffect(() => {
+    setUserMessages([
+      ...(datasourceState && datasourceState.state && datasource && activeDatasourceId
+        ? datasource.getUserMessages(datasourceState.state, {
+            frame: framePublicAPI,
+            setState: (newStateOrUpdater) => {
+              dispatch(
+                updateDatasourceState({
+                  newDatasourceState:
+                    typeof newStateOrUpdater === 'function'
+                      ? newStateOrUpdater(datasourceState.state)
+                      : newStateOrUpdater,
+                  datasourceId: activeDatasourceId,
+                })
+              );
+            },
+          })
+        : []),
+      ...(visualizationState?.activeId && visualizationState.state
+        ? visualization?.getUserMessages?.(visualizationState.state, {
+            frame: framePublicAPI,
+          }) ?? []
+        : []),
+      ...getApplicationUserMessages({
+        visualizationType,
+        visualization,
+        visualizationState,
+        activeDatasource: datasource,
+        activeDatasourceState: datasourceState,
+        core: coreStart,
+        dataViews: framePublicAPI.dataViews,
+      }),
+    ]);
+  }, [
+    activeDatasourceId,
+    datasource,
+    datasourceState,
+    dispatch,
+    framePublicAPI,
+    visualization,
+    visualizationState,
+    visualizationType,
+    coreStart,
+  ]);
+
+  const getUserMessages: UserMessagesGetter = (locationId, filterArgs) =>
+    filterAndSortUserMessages(
+      [...userMessages, ...Object.values(additionalUserMessages)],
+      locationId,
+      filterArgs ?? {}
+    );
+
+  const addUserMessages: AddUserMessages = (messages) => {
+    const newMessageMap = {
+      ...additionalUserMessages,
+    };
+
+    const addedMessageIds: string[] = [];
+    messages.forEach((message) => {
+      if (!newMessageMap[message.uniqueId]) {
+        addedMessageIds.push(message.uniqueId);
+        newMessageMap[message.uniqueId] = message;
+      }
+    });
+
+    if (addedMessageIds.length) {
+      setAdditionalUserMessages(newMessageMap);
+    }
+
+    return () => {
+      const withMessagesRemoved = {
+        ...additionalUserMessages,
+      };
+
+      addedMessageIds.forEach((id) => delete withMessagesRemoved[id]);
+
+      setAdditionalUserMessages(withMessagesRemoved);
+    };
+  };
+  return { getUserMessages, addUserMessages };
+};
--- a/x-pack/plugins/lens/public/datasources/form_based/form_based.test.ts
+++ b/x-pack/plugins/lens/public/datasources/form_based/form_based.test.ts
@ -26,7 +26,6 @@ import {
  Datasource,
  FramePublicAPI,
  OperationDescriptor,
-  FrameDatasourceAPI,
  UserMessage,
 } from '../../types';
 import { getFieldByNameFactory } from './pure_helpers';
@ -49,8 +48,9 @@ import {
 import { createMockedFullReference } from './operations/mocks';
 import { cloneDeep } from 'lodash';
 import { Datatable, DatatableColumn } from '@kbn/expressions-plugin/common';
-import { createMockFramePublicAPI } from '../../mocks';
 import { filterAndSortUserMessages } from '../../app_plugin/get_application_user_messages';
+import { createMockFramePublicAPI } from '../../mocks';
+import { createMockDataViewsState } from '../../data_views_service/mocks';

 jest.mock('./loader');
 jest.mock('../../id_generator');
@ -3062,22 +3062,6 @@ describe('IndexPattern Data Source', () => {
  });

  describe('#getUserMessages', () => {
-    function createMockFrameDatasourceAPI({
-      activeData,
-      dataViews,
-    }: Partial<Omit<FramePublicAPI, 'dataViews'>> & {
-      dataViews?: Partial<FramePublicAPI['dataViews']>;
-    }): FrameDatasourceAPI {
-      return {
-        ...createMockFramePublicAPI({
-          activeData,
-          dataViews,
-        }),
-        query: { query: '', language: 'kuery' },
-        filters: [],
-      };
-    }
-
    describe('error messages', () => {
      it('should generate error messages for a single layer', () => {
        (getErrorMessages as jest.Mock).mockClear();
@ -3094,7 +3078,9 @@ describe('IndexPattern Data Source', () => {
        };
        expect(
          FormBasedDatasource.getUserMessages(state, {
-            frame: createMockFrameDatasourceAPI({ dataViews: { indexPatterns } }),
+            frame: createMockFramePublicAPI({
+              dataViews: createMockDataViewsState({ indexPatterns }),
+            }),
            setState: () => {},
          })
        ).toMatchInlineSnapshot(`
@ -3146,7 +3132,9 @@ describe('IndexPattern Data Source', () => {
        };
        expect(
          FormBasedDatasource.getUserMessages(state, {
-            frame: createMockFrameDatasourceAPI({ dataViews: { indexPatterns } }),
+            frame: createMockFramePublicAPI({
+              dataViews: createMockDataViewsState({ indexPatterns }),
+            }),
            setState: () => {},
          })
        ).toMatchInlineSnapshot(`
@ -3235,7 +3223,9 @@ describe('IndexPattern Data Source', () => {
          (getErrorMessages as jest.Mock).mockReturnValueOnce([]);

          const messages = FormBasedDatasource.getUserMessages(state, {
-            frame: createMockFrameDatasourceAPI({ dataViews: { indexPatterns } }),
+            frame: createMockFramePublicAPI({
+              dataViews: createMockDataViewsState({ indexPatterns }),
+            }),
            setState: () => {},
          });

@ -3273,7 +3263,9 @@ describe('IndexPattern Data Source', () => {
          ] as ReturnType<typeof getErrorMessages>);

          const messages = FormBasedDatasource.getUserMessages(state, {
-            frame: createMockFrameDatasourceAPI({ dataViews: { indexPatterns } }),
+            frame: createMockFramePublicAPI({
+              dataViews: createMockDataViewsState({ indexPatterns }),
+            }),
            setState: () => {},
          });

@ -3303,7 +3295,7 @@ describe('IndexPattern Data Source', () => {

    describe('warning messages', () => {
      let state: FormBasedPrivateState;
-      let framePublicAPI: FrameDatasourceAPI;
+      let framePublicAPI: FramePublicAPI;

      beforeEach(() => {
        (getErrorMessages as jest.Mock).mockReturnValueOnce([]);
@ -3385,7 +3377,7 @@ describe('IndexPattern Data Source', () => {
          currentIndexPatternId: '1',
        };

-        framePublicAPI = createMockFrameDatasourceAPI({
+        framePublicAPI = createMockFramePublicAPI({
          activeData: {
            first: {
              type: 'datatable',
@ -3419,9 +3411,9 @@ describe('IndexPattern Data Source', () => {
              ],
            },
          },
-          dataViews: {
+          dataViews: createMockDataViewsState({
            indexPatterns: expectedIndexPatterns,
-          },
+          }),
        });
      });

@ -3549,13 +3541,13 @@ describe('IndexPattern Data Source', () => {
              currentIndexPatternId: '1',
            },
            {
-              frame: createMockFrameDatasourceAPI({
+              frame: createMockFramePublicAPI({
                activeData: {
                  first: createDatatableForLayer(0),
                },
-                dataViews: {
+                dataViews: createMockDataViewsState({
                  indexPatterns: expectedIndexPatterns,
-                },
+                }),
              }),
              setState: () => {},
              visualizationInfo: { layers: [] },
@ -3574,14 +3566,14 @@ describe('IndexPattern Data Source', () => {
          currentIndexPatternId: '1',
        };
        const messages = FormBasedDatasource.getUserMessages!(state, {
-          frame: createMockFrameDatasourceAPI({
+          frame: createMockFramePublicAPI({
            activeData: {
              first: createDatatableForLayer(0),
              second: createDatatableForLayer(1),
            },
-            dataViews: {
+            dataViews: createMockDataViewsState({
              indexPatterns: expectedIndexPatterns,
-            },
+            }),
          }),
          setState: () => {},
          visualizationInfo: { layers: [] },
--- a/x-pack/plugins/lens/public/datasources/form_based/form_based.tsx
+++ b/x-pack/plugins/lens/public/datasources/form_based/form_based.tsx
@ -38,7 +38,6 @@ import type {
  IndexPatternRef,
  DataSourceInfo,
  UserMessage,
-  FrameDatasourceAPI,
  StateSetter,
  IndexPatternMap,
 } from '../../types';
@ -749,18 +748,12 @@ export function getFormBasedDatasource({
    getDatasourceSuggestionsForVisualizeField,
    getDatasourceSuggestionsForVisualizeCharts,

-    getUserMessages(state, { frame: frameDatasourceAPI, setState, visualizationInfo }) {
+    getUserMessages(state, { frame: framePublicAPI, setState, visualizationInfo }) {
      if (!state) {
        return [];
      }

-      const layerErrorMessages = getLayerErrorMessages(
-        state,
-        frameDatasourceAPI,
-        setState,
-        core,
-        data
-      );
+      const layerErrorMessages = getLayerErrorMessages(state, framePublicAPI, setState, core, data);

      const dimensionErrorMessages = getInvalidDimensionErrorMessages(
        state,
@ -770,8 +763,8 @@ export function getFormBasedDatasource({
          return !isColumnInvalid(
            layer,
            columnId,
-            frameDatasourceAPI.dataViews.indexPatterns[layer.indexPatternId],
-            frameDatasourceAPI.dateRange,
+            framePublicAPI.dataViews.indexPatterns[layer.indexPatternId],
+            framePublicAPI.dateRange,
            uiSettings.get(UI_SETTINGS.HISTOGRAM_BAR_TARGET)
          );
        }
@ -779,11 +772,8 @@ export function getFormBasedDatasource({

      const warningMessages = [
        ...[
-          ...(getStateTimeShiftWarningMessages(
-            data.datatableUtilities,
-            state,
-            frameDatasourceAPI
-          ) || []),
+          ...(getStateTimeShiftWarningMessages(data.datatableUtilities, state, framePublicAPI) ||
+            []),
        ].map((longMessage) => {
          const message: UserMessage = {
            severity: 'warning',
@ -798,14 +788,14 @@ export function getFormBasedDatasource({
        ...getPrecisionErrorWarningMessages(
          data.datatableUtilities,
          state,
-          frameDatasourceAPI,
+          framePublicAPI,
          core.docLinks,
          setState
        ),
-        ...getUnsupportedOperationsWarningMessage(state, frameDatasourceAPI, core.docLinks),
+        ...getUnsupportedOperationsWarningMessage(state, framePublicAPI, core.docLinks),
      ];

-      const infoMessages = getNotifiableFeatures(state, frameDatasourceAPI, visualizationInfo);
+      const infoMessages = getNotifiableFeatures(state, framePublicAPI, visualizationInfo);

      return layerErrorMessages.concat(dimensionErrorMessages, warningMessages, infoMessages);
    },
@ -922,12 +912,12 @@ function blankLayer(indexPatternId: string, linkToLayers?: string[]): FormBasedL

 function getLayerErrorMessages(
  state: FormBasedPrivateState,
-  frameDatasourceAPI: FrameDatasourceAPI,
+  framePublicAPI: FramePublicAPI,
  setState: StateSetter<FormBasedPrivateState, unknown>,
  core: CoreStart,
  data: DataPublicPluginStart
 ) {
-  const indexPatterns = frameDatasourceAPI.dataViews.indexPatterns;
+  const indexPatterns = framePublicAPI.dataViews.indexPatterns;

  const layerErrors: UserMessage[][] = Object.entries(state.layers)
    .filter(([_, layer]) => !!indexPatterns[layer.indexPatternId])
@ -954,7 +944,7 @@ function getLayerErrorMessages(
                  <EuiButton
                    data-test-subj="errorFixAction"
                    onClick={async () => {
-                      const newState = await error.fixAction?.newState(frameDatasourceAPI);
+                      const newState = await error.fixAction?.newState(framePublicAPI);
                      if (newState) {
                        setState(newState);
                      }
--- a/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/index.ts
+++ b/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/index.ts
@ -51,7 +51,7 @@ import {
 import { staticValueOperation } from './static_value';
 import { lastValueOperation } from './last_value';
 import type {
-  FrameDatasourceAPI,
+  FramePublicAPI,
  IndexPattern,
  IndexPatternField,
  OperationMetadata,
@ -477,7 +477,7 @@ export type FieldBasedOperationErrorMessage =
        newState: (
          data: DataPublicPluginStart,
          core: CoreStart,
-          frame: FrameDatasourceAPI,
+          frame: FramePublicAPI,
          layerId: string
        ) => Promise<FormBasedLayer>;
      };
--- a/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/terms/helpers.test.ts
+++ b/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/terms/helpers.test.ts
@ -7,7 +7,7 @@

 import { dataPluginMock } from '@kbn/data-plugin/public/mocks';
 import { coreMock as corePluginMock } from '@kbn/core/public/mocks';
-import type { FrameDatasourceAPI } from '../../../../../types';
+import type { FramePublicAPI } from '../../../../../types';
 import type { CountIndexPatternColumn } from '..';
 import type { TermsIndexPatternColumn } from './types';
 import type { GenericIndexPatternColumn } from '../../../form_based';
@ -245,7 +245,7 @@ describe('getDisallowedTermsMessage()', () => {
          fromDate: '2020',
          toDate: '2021',
        },
-      } as unknown as FrameDatasourceAPI,
+      } as unknown as FramePublicAPI,
      'first'
    );

@ -299,7 +299,7 @@ describe('getDisallowedTermsMessage()', () => {
            rows: [{ col1: 'myTerm' }, { col1: 'myOtherTerm' }],
          },
        },
-      } as unknown as FrameDatasourceAPI,
+      } as unknown as FramePublicAPI,
      'first'
    );

@ -335,7 +335,7 @@ describe('getDisallowedTermsMessage()', () => {
          fromDate: '2020',
          toDate: '2021',
        },
-      } as unknown as FrameDatasourceAPI,
+      } as unknown as FramePublicAPI,
      'first'
    );

@ -385,7 +385,7 @@ describe('getDisallowedTermsMessage()', () => {
            ],
          },
        },
-      } as unknown as FrameDatasourceAPI,
+      } as unknown as FramePublicAPI,
      'first'
    );

--- a/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/terms/helpers.ts
+++ b/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/terms/helpers.ts
@ -17,7 +17,7 @@ import { GenericIndexPatternColumn, operationDefinitionMap } from '..';
 import { defaultLabel } from '../filters';
 import { isReferenced } from '../../layer_helpers';

-import type { FrameDatasourceAPI, IndexPattern, IndexPatternField } from '../../../../../types';
+import type { FramePublicAPI, IndexPattern, IndexPatternField } from '../../../../../types';
 import type { FiltersIndexPatternColumn } from '..';
 import type { TermsIndexPatternColumn } from './types';
 import type { LastValueIndexPatternColumn } from '../last_value';
@ -126,7 +126,7 @@ export function getDisallowedTermsMessage(
      newState: async (
        data: DataPublicPluginStart,
        core: CoreStart,
-        frame: FrameDatasourceAPI,
+        frame: FramePublicAPI,
        layerId: string
      ) => {
        const currentColumn = layer.columns[columnId] as TermsIndexPatternColumn;
--- a/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/terms/terms.test.tsx
+++ b/x-pack/plugins/lens/public/datasources/form_based/operations/definitions/terms/terms.test.tsx
@ -33,7 +33,7 @@ import {
  operationDefinitionMap,
 } from '..';
 import { FormBasedLayer, FormBasedPrivateState } from '../../../types';
-import { FrameDatasourceAPI } from '../../../../../types';
+import { FramePublicAPI } from '../../../../../types';
 import { DateHistogramIndexPatternColumn } from '../date_histogram';
 import { getOperationSupportMatrix } from '../../../dimension_panel/operation_support';
 import { FieldSelect } from '../../../dimension_panel/field_select';
@ -2867,7 +2867,7 @@ describe('terms', () => {
              fromDate: '2020',
              toDate: '2021',
            },
-          } as unknown as FrameDatasourceAPI,
+          } as unknown as FramePublicAPI,
          'first'
        );
        expect(newLayer.columns.col1).toEqual(
--- a/x-pack/plugins/lens/public/datasources/form_based/operations/layer_helpers.ts
+++ b/x-pack/plugins/lens/public/datasources/form_based/operations/layer_helpers.ts
@ -13,7 +13,7 @@ import { DataPublicPluginStart, UI_SETTINGS } from '@kbn/data-plugin/public';
 import type { DateRange } from '../../../../common/types';
 import type {
  DatasourceFixAction,
-  FrameDatasourceAPI,
+  FramePublicAPI,
  IndexPattern,
  IndexPatternField,
  OperationMetadata,
@ -1594,7 +1594,7 @@ export function getErrorMessages(
        fixAction: errorMessage.fixAction
          ? {
              ...errorMessage.fixAction,
-              newState: async (frame: FrameDatasourceAPI) => ({
+              newState: async (frame: FramePublicAPI) => ({
                ...state,
                layers: {
                  ...state.layers,
--- a/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts
+++ b/x-pack/plugins/lens/public/datasources/text_based/text_based_languages.test.ts
@ -13,7 +13,7 @@ import { dataPluginMock } from '@kbn/data-plugin/public/mocks';
 import { dataViewPluginMocks } from '@kbn/data-views-plugin/public/mocks';
 import { getTextBasedDatasource } from './text_based_languages';
 import { generateId } from '../../id_generator';
-import { DatasourcePublicAPI, Datasource, FrameDatasourceAPI } from '../../types';
+import { DatasourcePublicAPI, Datasource, FramePublicAPI } from '../../types';

 jest.mock('../../id_generator');

@ -551,7 +551,7 @@ describe('Textbased Data Source', () => {
      } as unknown as TextBasedPrivateState;
      expect(
        TextBasedDatasource.getUserMessages(state, {
-          frame: { dataViews: indexPatterns } as unknown as FrameDatasourceAPI,
+          frame: { dataViews: indexPatterns } as unknown as FramePublicAPI,
          setState: () => {},
        })
      ).toMatchInlineSnapshot(`
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/config_panel.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/config_panel/config_panel.test.tsx
@ -477,12 +477,14 @@ describe('ConfigPanel', () => {
      expect(visualizationMap.testVis.setDimension).toHaveBeenCalledWith({
        columnId: 'newId',
        frame: {
+          dataViews: expect.anything(),
          activeData: undefined,
          datasourceLayers: {
            a: expect.anything(),
          },
          dateRange: expect.anything(),
-          dataViews: expect.anything(),
+          filters: [],
+          query: undefined,
        },
        groupId: 'a',
        layerId: 'newId',
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/editor_frame.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/editor_frame.test.tsx
@ -38,6 +38,7 @@ import {
  DatasourceMock,
  createExpressionRendererMock,
  mockStoreDeps,
+  renderWithReduxStore,
 } from '../../mocks';
 import { inspectorPluginMock } from '@kbn/inspector-plugin/public/mocks';
 import { ReactExpressionRendererType } from '@kbn/expressions-plugin/public';
@ -283,7 +284,7 @@ describe('editor_frame', () => {
      const props = {
        ...getDefaultProps(),
        visualizationMap: {
-          testVis: mockVisualization,
+          testVis: { ...mockVisualization, toExpression: () => null },
        },
        datasourceMap: {
          testDatasource: mockDatasource,
@ -291,18 +292,23 @@ describe('editor_frame', () => {

        ExpressionRenderer: expressionRendererMock,
      };
-      await mountWithProvider(<EditorFrame {...props} />, {
-        preloadedState: {
-          activeDatasourceId: 'testDatasource',
-          visualization: { activeId: mockVisualization.id, state: {} },
-          datasourceStates: {
-            testDatasource: {
-              isLoading: false,
-              state: '',
+      renderWithReduxStore(
+        <EditorFrame {...props} />,
+        {},
+        {
+          preloadedState: {
+            activeDatasourceId: 'testDatasource',
+            visualization: { activeId: mockVisualization.id, state: {} },
+            datasourceStates: {
+              testDatasource: {
+                isLoading: false,
+                state: '',
+              },
            },
          },
-        },
-      });
+        }
+      );
+
      const updatedState = {};
      const setDatasourceState = (mockDatasource.DataPanelComponent as jest.Mock).mock.calls[0][0]
        .setState;
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/state_helpers.ts
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/state_helpers.ts
@ -35,7 +35,7 @@ import type {
 import { buildExpression } from './expression_helpers';
 import { Document } from '../../persistence/saved_object_store';
 import { getActiveDatasourceIdFromDoc, sortDataViewRefs } from '../../utils';
-import type { DatasourceStates, VisualizationState } from '../../state_management';
+import type { DatasourceState, DatasourceStates, VisualizationState } from '../../state_management';
 import { readFromStorage } from '../../settings_storage';
 import { loadIndexPatternRefs, loadIndexPatterns } from '../../data_views_service/loader';
 import { getDatasourceLayers } from '../../state_management/utils';
@ -461,7 +461,7 @@ export async function persistedStateToExpression(

 export function getMissingIndexPattern(
  currentDatasource: Datasource | null | undefined,
-  currentDatasourceState: { isLoading: boolean; state: unknown } | null,
+  currentDatasourceState: DatasourceState | null,
  indexPatterns: IndexPatternMap
 ) {
  if (
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/suggestion_panel.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/suggestion_panel.test.tsx
@ -381,8 +381,9 @@ describe('suggestion_panel', () => {
      },
    ] as Suggestion[]);

-    (mockVisualization.toPreviewExpression as jest.Mock).mockReturnValueOnce(undefined);
-    (mockVisualization.toPreviewExpression as jest.Mock).mockReturnValueOnce('test | expression');
+    (mockVisualization.toPreviewExpression as jest.Mock)
+      .mockReturnValue(undefined)
+      .mockReturnValueOnce('test | expression');
    mockDatasource.toExpression.mockReturnValue('datasource_expression');

    mountWithProvider(<SuggestionPanel {...defaultProps} frame={createMockFramePublicAPI()} />);
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/suggestion_panel.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/suggestion_panel.tsx
@ -41,7 +41,6 @@ import {
  VisualizationMap,
  DatasourceLayers,
  UserMessagesGetter,
-  FrameDatasourceAPI,
 } from '../../types';
 import { getSuggestions, switchToSuggestion } from './suggestion_helpers';
 import { getDatasourceExpressionsByLayers } from './expression_helpers';
@ -63,7 +62,7 @@ import {
  selectChangesApplied,
  applyChanges,
  selectStagedActiveData,
-  selectFrameDatasourceAPI,
+  selectFramePublicAPI,
 } from '../../state_management';
 import { filterAndSortUserMessages } from '../../app_plugin/get_application_user_messages';
 const MAX_SUGGESTIONS_DISPLAYED = 5;
@ -74,7 +73,7 @@ const configurationsValid = (
  currentDatasourceState: unknown,
  currentVisualization: Visualization,
  currentVisualizationState: unknown,
-  frame: FrameDatasourceAPI
+  frame: FramePublicAPI
 ): boolean => {
  try {
    return (
@ -241,9 +240,7 @@ export function SuggestionPanel({
  const currentVisualization = useLensSelector(selectCurrentVisualization);
  const currentDatasourceStates = useLensSelector(selectCurrentDatasourceStates);

-  const frameDatasourceAPI = useLensSelector((state) =>
-    selectFrameDatasourceAPI(state, datasourceMap)
-  );
+  const framePublicAPI = useLensSelector((state) => selectFramePublicAPI(state, datasourceMap));
  const changesApplied = useLensSelector(selectChangesApplied);
  // get user's selection from localStorage, this key defines if the suggestions panel will be hidden or not
  const [hideSuggestions, setHideSuggestions] = useLocalStorage(
@ -289,7 +286,7 @@ export function SuggestionPanel({
                  suggestionDatasourceState,
                  visualizationMap[visualizationId],
                  suggestionVisualizationState,
-                  frameDatasourceAPI
+                  framePublicAPI
                )
              );
            }
--- a/x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/workspace_panel.test.tsx
+++ b/x-pack/plugins/lens/public/editor_frame_service/editor_frame/workspace_panel/workspace_panel.test.tsx
@ -380,7 +380,7 @@ describe('workspace_panel', () => {
          }}
          framePublicAPI={framePublicAPI}
          visualizationMap={{
-            testVis: mockVisualization,
+            testVis: { ...mockVisualization, toExpression: () => null },
          }}
          ExpressionRenderer={expressionRendererMock}
        />,
--- a/x-pack/plugins/lens/public/embeddable/embeddable.tsx
+++ b/x-pack/plugins/lens/public/embeddable/embeddable.tsx
@ -97,7 +97,7 @@ import {
  GetCompatibleCellValueActions,
  UserMessage,
  IndexPatternRef,
-  FrameDatasourceAPI,
+  FramePublicAPI,
  AddUserMessages,
  isMessageRemovable,
  UserMessagesGetter,
@ -608,11 +608,14 @@ export class Embeddable
    userMessages.push(
      ...getApplicationUserMessages({
        visualizationType: this.savedVis?.visualizationType,
-        visualization: {
+        visualizationState: {
          state: this.activeVisualizationState,
          activeId: this.activeVisualizationId,
        },
-        visualizationMap: this.deps.visualizationMap,
+        visualization:
+          this.activeVisualizationId && this.deps.visualizationMap[this.activeVisualizationId]
+            ? this.deps.visualizationMap[this.activeVisualizationId]
+            : undefined,
        activeDatasource: this.activeDatasource,
        activeDatasourceState: {
          isLoading: !this.activeDatasourceState,
@ -631,7 +634,7 @@ export class Embeddable
    }
    const mergedSearchContext = this.getMergedSearchContext();

-    const frameDatasourceAPI: FrameDatasourceAPI = {
+    const framePublicAPI: FramePublicAPI = {
      dataViews: {
        indexPatterns: this.indexPatterns,
        indexPatternRefs: this.indexPatternRefs,
@ -658,14 +661,14 @@ export class Embeddable
    userMessages.push(
      ...(this.activeDatasource?.getUserMessages(this.activeDatasourceState, {
        setState: () => {},
-        frame: frameDatasourceAPI,
+        frame: framePublicAPI,
        visualizationInfo: this.activeVisualization?.getVisualizationInfo?.(
          this.activeVisualizationState,
-          frameDatasourceAPI
+          framePublicAPI
        ),
      }) ?? []),
      ...(this.activeVisualization?.getUserMessages?.(this.activeVisualizationState, {
-        frame: frameDatasourceAPI,
+        frame: framePublicAPI,
      }) ?? [])
    );

--- a/x-pack/plugins/lens/public/mocks/datasource_mock.tsx
+++ b/x-pack/plugins/lens/public/mocks/datasource_mock.tsx
@ -44,7 +44,9 @@ export function createMockDatasource(
    getRenderEventCounters: jest.fn((_state) => []),
    getPublicAPI: jest.fn().mockReturnValue(publicAPIMock),
    initialize: jest.fn((_state?) => {}),
-    toExpression: jest.fn((_frame, _state, _indexPatterns, dateRange, nowInstant) => null),
+    toExpression: jest.fn(
+      (_frame, _state, _indexPatterns, dateRange, nowInstant) => 'datasource_expression'
+    ),
    insertLayer: jest.fn((_state, _newLayerId) => ({})),
    removeLayer: jest.fn((state, layerId) => ({ newState: state, removedLayerIds: [layerId] })),
    cloneLayer: jest.fn((_state, _layerId, _newLayerId, getNewId) => {}),
@ -62,7 +64,7 @@ export function createMockDatasource(
    getUserMessages: jest.fn((_state, _deps) => []),
    checkIntegrity: jest.fn((_state, _indexPatterns) => []),
    isTimeBased: jest.fn(),
-    isEqual: jest.fn(),
+    isEqual: jest.fn((a, b, c, d) => a === c),
    getUsedDataView: jest.fn((state, layer) => 'mockip'),
    getUsedDataViews: jest.fn(),
    onRefreshIndexPattern: jest.fn(),
--- a/x-pack/plugins/lens/public/mocks/index.ts
+++ b/x-pack/plugins/lens/public/mocks/index.ts
@ -8,7 +8,7 @@
 import { DragContextState, DragContextValue } from '@kbn/dom-drag-drop';
 import { DatatableColumnType } from '@kbn/expressions-plugin/common';
 import { createMockDataViewsState } from '../data_views_service/mocks';
-import { FramePublicAPI, FrameDatasourceAPI } from '../types';
+import { FramePublicAPI } from '../types';
 export { mockDataPlugin } from './data_plugin_mock';
 export {
  visualizationMap,
@ -32,40 +32,16 @@ export { lensPluginMock } from './lens_plugin_mock';

 export type FrameMock = jest.Mocked<FramePublicAPI>;

-export const createMockFramePublicAPI = ({
-  datasourceLayers,
-  dateRange,
-  dataViews,
-  activeData,
-}: Partial<Omit<FramePublicAPI, 'dataViews'>> & {
-  dataViews?: Partial<FramePublicAPI['dataViews']>;
-} = {}): FrameMock => ({
-  datasourceLayers: datasourceLayers ?? {},
-  dateRange: dateRange ?? {
+export const createMockFramePublicAPI = (overrides: Partial<FramePublicAPI> = {}): FrameMock => ({
+  datasourceLayers: {},
+  dateRange: {
    fromDate: '2022-03-17T08:25:00.000Z',
    toDate: '2022-04-17T08:25:00.000Z',
  },
-  dataViews: createMockDataViewsState(dataViews),
-  activeData,
-});
-
-export type FrameDatasourceMock = jest.Mocked<FrameDatasourceAPI>;
-
-export const createMockFrameDatasourceAPI = ({
-  datasourceLayers,
-  dateRange,
-  dataViews,
-  query,
-  filters,
-}: Partial<FrameDatasourceAPI> = {}): FrameDatasourceMock => ({
-  datasourceLayers: datasourceLayers ?? {},
-  dateRange: dateRange ?? {
-    fromDate: '2022-03-17T08:25:00.000Z',
-    toDate: '2022-04-17T08:25:00.000Z',
-  },
-  query: query ?? { query: '', language: 'lucene' },
-  filters: filters ?? [],
-  dataViews: createMockDataViewsState(dataViews),
+  dataViews: createMockDataViewsState(),
+  query: { query: '', language: 'lucene' },
+  filters: [],
+  ...overrides,
 });

 export function createMockedDragDropContext(
--- a/x-pack/plugins/lens/public/mocks/visualization_mock.tsx
+++ b/x-pack/plugins/lens/public/mocks/visualization_mock.tsx
@ -44,9 +44,9 @@ export function createMockVisualization(id = 'testVis'): jest.Mocked<Visualizati
        },
      ],
    })),
-    toExpression: jest.fn((_state, _frame) => null),
-    toPreviewExpression: jest.fn((_state, _frame) => null),
-
+    toExpression: jest.fn((_state, _frame) => 'expression'),
+    toPreviewExpression: jest.fn((_state, _frame) => 'expression'),
+    getUserMessages: jest.fn((_state) => []),
    setDimension: jest.fn(),
    removeDimension: jest.fn(),
    DimensionEditorComponent: jest.fn(() => <div />),
--- a/x-pack/plugins/lens/public/state_management/selectors.ts
+++ b/x-pack/plugins/lens/public/state_management/selectors.ts
@ -220,10 +220,10 @@ export const selectFramePublicAPI = createSelector(
    selectCurrentDatasourceStates,
    selectActiveData,
    selectInjectedDependencies as SelectInjectedDependenciesFunction<DatasourceMap>,
-    selectResolvedDateRange,
    selectDataViews,
+    selectExecutionContext,
  ],
-  (datasourceStates, activeData, datasourceMap, dateRange, dataViews) => {
+  (datasourceStates, activeData, datasourceMap, dataViews, context) => {
    return {
      datasourceLayers: getDatasourceLayers(
        datasourceStates,
@ -231,13 +231,8 @@ export const selectFramePublicAPI = createSelector(
        dataViews.indexPatterns
      ),
      activeData,
-      dateRange,
      dataViews,
+      ...context,
    };
  }
 );
-
-export const selectFrameDatasourceAPI = createSelector(
-  [selectFramePublicAPI, selectExecutionContext],
-  (framePublicAPI, context) => ({ ...context, ...framePublicAPI })
-);
--- a/x-pack/plugins/lens/public/state_management/types.ts
+++ b/x-pack/plugins/lens/public/state_management/types.ts
@ -34,7 +34,12 @@ export interface DataViewsState {
  indexPatterns: Record<string, IndexPattern>;
 }

-export type DatasourceStates = Record<string, { isLoading: boolean; state: unknown }>;
+export interface DatasourceState {
+  isLoading: boolean;
+  state: unknown;
+}
+
+export type DatasourceStates = Record<string, DatasourceState>;
 export interface PreviewState {
  visualization: VisualizationState;
  datasourceStates: DatasourceStates;
--- a/x-pack/plugins/lens/public/types.ts
+++ b/x-pack/plugins/lens/public/types.ts
@ -460,7 +460,7 @@ export interface Datasource<T = unknown, P = unknown> {
  getUserMessages: (
    state: T,
    deps: {
-      frame: FrameDatasourceAPI;
+      frame: FramePublicAPI;
      setState: StateSetter<T>;
      visualizationInfo?: VisualizationInfo;
    }
@ -513,7 +513,7 @@ export interface Datasource<T = unknown, P = unknown> {

 export interface DatasourceFixAction<T> {
  label: string;
-  newState: (frame: FrameDatasourceAPI) => Promise<T>;
+  newState: (frame: FramePublicAPI) => Promise<T>;
 }

 /**
@ -925,6 +925,8 @@ export interface VisualizationSuggestion<T = unknown> {
 export type DatasourceLayers = Partial<Record<string, DatasourcePublicAPI>>;

 export interface FramePublicAPI {
+  query: Query;
+  filters: Filter[];
  datasourceLayers: DatasourceLayers;
  dateRange: DateRange;
  /**
@ -936,11 +938,6 @@ export interface FramePublicAPI {
  dataViews: DataViewsState;
 }

-export interface FrameDatasourceAPI extends FramePublicAPI {
-  query: Query;
-  filters: Filter[];
-}
-
 /**
 * A visualization type advertised to the user in the chart switcher
 */
--- a/x-pack/plugins/lens/public/visualizations/xy/snapshots/to_expression.test.ts.snap
+++ b/x-pack/plugins/lens/public/visualizations/xy/snapshots/to_expression.test.ts.snap
@ -27,6 +27,11 @@ Object {
        "layers": Array [
          Object {
            "chain": Array [
+              Object {
+                "arguments": Object {},
+                "function": "datasource_expression",
+                "type": "function",
+              },
              Object {
                "arguments": Object {
                  "accessors": Array [
--- a/x-pack/plugins/lens/public/visualizations/xy/to_expression.test.ts
+++ b/x-pack/plugins/lens/public/visualizations/xy/to_expression.test.ts
@ -265,7 +265,7 @@ describe('#toExpression', () => {
    expect(mockDatasource.publicAPIMock.getOperationForColumnId).toHaveBeenCalledWith('c');
    expect(mockDatasource.publicAPIMock.getOperationForColumnId).toHaveBeenCalledWith('d');
    expect(
-      (expression.chain[0].arguments.layers[0] as Ast).chain[0].arguments.columnToLabel
+      (expression.chain[0].arguments.layers[0] as Ast).chain[1].arguments.columnToLabel
    ).toEqual([
      JSON.stringify({
        b: 'col_b',
@ -536,13 +536,18 @@ describe('#toExpression', () => {
      datasourceExpressionsByLayers
    ) as Ast;

-    function getYConfigColorForLayer(ast: Ast, index: number) {
+    function getYConfigColorForDataLayer(ast: Ast, index: number) {
+      return (
+        (ast.chain[0].arguments.layers[index] as Ast).chain[1].arguments.decorations[0] as Ast
+      ).chain[0].arguments?.color;
+    }
+    function getYConfigColorForReferenceLayer(ast: Ast, index: number) {
      return (
        (ast.chain[0].arguments.layers[index] as Ast).chain[0].arguments.decorations[0] as Ast
-      ).chain[0].arguments.color;
+      ).chain[0].arguments?.color;
    }
-    expect(getYConfigColorForLayer(expression, 0)).toBeUndefined();
-    expect(getYConfigColorForLayer(expression, 1)).toEqual([defaultReferenceLineColor]);
+    expect(getYConfigColorForDataLayer(expression, 0)).toBeUndefined();
+    expect(getYConfigColorForReferenceLayer(expression, 1)).toEqual([defaultReferenceLineColor]);
  });

  it('should ignore annotation layers with no event configured', () => {