Refactor attribute service (#78414) (#79621)

* Making saveMethod mandatory in attribute service

* Making unwrap method mandatory

* Making book embeddable respect new attribute service

* Remove savedObjectsClient from attribute service

* Add checkForDuplicateTitle method to book embeddable

* Make options mandatory on attribute service

* Changing Lens attribute service

* Somw more typescript fixes

* Fixing attribute service typescript and tests

* Fixing typescript errors

* Unsetting feature flag

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

examples/embeddable_examples/kibana.json

@@ -4,7 +4,7 @@
   "kibanaVersion": "kibana",
   "server": true,
   "ui": true,
-  "requiredPlugins": ["embeddable", "uiActions", "dashboard"],
+  "requiredPlugins": ["embeddable", "uiActions", "dashboard", "savedObjects"],
   "optionalPlugins": [],
   "extraPublicDirs": ["public/todo", "public/hello_world", "public/todo/todo_ref_embeddable"],
   "requiredBundles": ["kibanaReact"]

examples/embeddable_examples/public/book/book_embeddable_factory.tsx

@@ -33,12 +33,19 @@ import {
   BookEmbeddableOutput,
 } from './book_embeddable';
 import { CreateEditBookComponent } from './create_edit_book_component';
-import { OverlayStart } from '../../../../src/core/public';
+import {
+  OverlayStart,
+  SavedObjectsClientContract,
+  SimpleSavedObject,
+} from '../../../../src/core/public';
 import { DashboardStart, AttributeService } from '../../../../src/plugins/dashboard/public';
+import { checkForDuplicateTitle, OnSaveProps } from '../../../../src/plugins/saved_objects/public';
 
 interface StartServices {
   getAttributeService: DashboardStart['getAttributeService'];
   openModal: OverlayStart['openModal'];
+  savedObjectsClient: SavedObjectsClientContract;
+  overlays: OverlayStart;
 }
 
 export type BookEmbeddableFactory = EmbeddableFactory<
@@ -117,11 +124,55 @@ export class BookEmbeddableFactoryDefinition
     });
   }
 
+  private async unwrapMethod(savedObjectId: string): Promise<BookSavedObjectAttributes> {
+    const { savedObjectsClient } = await this.getStartServices();
+    const savedObject: SimpleSavedObject<BookSavedObjectAttributes> = await savedObjectsClient.get<
+      BookSavedObjectAttributes
+    >(this.type, savedObjectId);
+    return { ...savedObject.attributes };
+  }
+
+  private async saveMethod(
+    type: string,
+    attributes: BookSavedObjectAttributes,
+    savedObjectId?: string
+  ) {
+    const { savedObjectsClient } = await this.getStartServices();
+    if (savedObjectId) {
+      return savedObjectsClient.update(type, savedObjectId, attributes);
+    }
+    return savedObjectsClient.create(type, attributes);
+  }
+
+  private async checkForDuplicateTitleMethod(props: OnSaveProps): Promise<true> {
+    const start = await this.getStartServices();
+    const { savedObjectsClient, overlays } = start;
+    return checkForDuplicateTitle(
+      {
+        title: props.newTitle,
+        copyOnSave: false,
+        lastSavedTitle: '',
+        getEsType: () => this.type,
+        getDisplayName: this.getDisplayName || (() => this.type),
+      },
+      props.isTitleDuplicateConfirmed,
+      props.onTitleDuplicate,
+      {
+        savedObjectsClient,
+        overlays,
+      }
+    );
+  }
+
   private async getAttributeService() {
     if (!this.attributeService) {
-      this.attributeService = await (await this.getStartServices()).getAttributeService<
+      this.attributeService = (await this.getStartServices()).getAttributeService<
         BookSavedObjectAttributes
-      >(this.type);
+      >(this.type, {
+        saveMethod: this.saveMethod.bind(this),
+        unwrapMethod: this.unwrapMethod.bind(this),
+        checkForDuplicateTitle: this.checkForDuplicateTitleMethod.bind(this),
+      });
     }
     return this.attributeService!;
   }

examples/embeddable_examples/public/book/edit_book_action.tsx

@@ -31,10 +31,13 @@ import {
 } from './book_embeddable';
 import { CreateEditBookComponent } from './create_edit_book_component';
 import { DashboardStart } from '../../../../src/plugins/dashboard/public';
+import { OnSaveProps } from '../../../../src/plugins/saved_objects/public';
+import { SavedObjectsClientContract } from '../../../../src/core/target/types/public/saved_objects';
 
 interface StartServices {
   openModal: OverlayStart['openModal'];
   getAttributeService: DashboardStart['getAttributeService'];
+  savedObjectsClient: SavedObjectsClientContract;
 }
 
 interface ActionContext {
@@ -56,8 +59,24 @@ export const createEditBookAction = (getStartServices: () => Promise<StartServic
       );
     },
     execute: async ({ embeddable }: ActionContext) => {
-      const { openModal, getAttributeService } = await getStartServices();
-      const attributeService = getAttributeService<BookSavedObjectAttributes>(BOOK_SAVED_OBJECT);
+      const { openModal, getAttributeService, savedObjectsClient } = await getStartServices();
+      const attributeService = getAttributeService<BookSavedObjectAttributes>(BOOK_SAVED_OBJECT, {
+        saveMethod: async (
+          type: string,
+          attributes: BookSavedObjectAttributes,
+          savedObjectId?: string
+        ) => {
+          if (savedObjectId) {
+            return savedObjectsClient.update(type, savedObjectId, attributes);
+          }
+          return savedObjectsClient.create(type, attributes);
+        },
+        checkForDuplicateTitle: (props: OnSaveProps) => {
+          return new Promise(() => {
+            return true;
+          });
+        },
+      });
       const onSave = async (attributes: BookSavedObjectAttributes, useRefType: boolean) => {
         const newInput = await attributeService.wrapAttributes(
           attributes,

examples/embeddable_examples/public/plugin.ts

@@ -22,7 +22,7 @@ import {
   EmbeddableStart,
   CONTEXT_MENU_TRIGGER,
 } from '../../../src/plugins/embeddable/public';
-import { Plugin, CoreSetup, CoreStart } from '../../../src/core/public';
+import { Plugin, CoreSetup, CoreStart, SavedObjectsClient } from '../../../src/core/public';
 import {
   HelloWorldEmbeddableFactory,
   HELLO_WORLD_EMBEDDABLE,
@@ -76,6 +76,7 @@ export interface EmbeddableExamplesSetupDependencies {
 export interface EmbeddableExamplesStartDependencies {
   embeddable: EmbeddableStart;
   dashboard: DashboardStart;
+  savedObjectsClient: SavedObjectsClient;
 }
 
 interface ExampleEmbeddableFactories {
@@ -158,12 +159,15 @@ export class EmbeddableExamplesPlugin
       new BookEmbeddableFactoryDefinition(async () => ({
         getAttributeService: (await core.getStartServices())[1].dashboard.getAttributeService,
         openModal: (await core.getStartServices())[0].overlays.openModal,
+        savedObjectsClient: (await core.getStartServices())[0].savedObjects.client,
+        overlays: (await core.getStartServices())[0].overlays,
       }))
     );
 
     const editBookAction = createEditBookAction(async () => ({
       getAttributeService: (await core.getStartServices())[1].dashboard.getAttributeService,
       openModal: (await core.getStartServices())[0].overlays.openModal,
+      savedObjectsClient: (await core.getStartServices())[0].savedObjects.client,
     }));
     deps.uiActions.registerAction(editBookAction);
     deps.uiActions.attachAction(CONTEXT_MENU_TRIGGER, editBookAction.id);

src/plugins/dashboard/public/attribute_service/attribute_service.mock.tsx

@@ -31,19 +31,16 @@ export const mockAttributeService = <
   R extends SavedObjectEmbeddableInput = SavedObjectEmbeddableInput
 >(
   type: string,
-  options?: AttributeServiceOptions<A>,
+  options: AttributeServiceOptions<A>,
   customCore?: jest.Mocked<CoreStart>
 ): AttributeService<A, V, R> => {
   const core = customCore ? customCore : coreMock.createStart();
-  const service = new AttributeService<A, V, R>(
+  return new AttributeService<A, V, R>(
     type,
     jest.fn(),
-    core.savedObjects.client,
-    core.overlays,
     core.i18n.Context,
     core.notifications.toasts,
-    jest.fn().mockReturnValue(() => ({ getDisplayName: () => type })),
-    options
+    options,
+    jest.fn().mockReturnValue(() => ({ getDisplayName: () => type }))
   );
-  return service;
 };

src/plugins/dashboard/public/attribute_service/attribute_service.test.ts

@@ -20,6 +20,7 @@
 import { ATTRIBUTE_SERVICE_KEY } from './attribute_service';
 import { mockAttributeService } from './attribute_service.mock';
 import { coreMock } from '../../../../core/public/mocks';
+import { OnSaveProps } from '../../../saved_objects/public/save_modal';
 
 interface TestAttributes {
   title: string;
@@ -37,6 +38,30 @@ describe('attributeService', () => {
   let attributes: TestAttributes;
   let byValueInput: TestByValueInput;
   let byReferenceInput: { id: string; savedObjectId: string };
+  const defaultSaveMethod = (
+    type: string,
+    testAttributes: TestAttributes,
+    savedObjectId?: string
+  ): Promise<{ id: string }> => {
+    return new Promise(() => {
+      return { id: '123' };
+    });
+  };
+  const defaultUnwrapMethod = (savedObjectId: string): Promise<TestAttributes> => {
+    return new Promise(() => {
+      return { ...attributes };
+    });
+  };
+  const defaultCheckForDuplicateTitle = (props: OnSaveProps): Promise<true> => {
+    return new Promise(() => {
+      return true;
+    });
+  };
+  const options = {
+    saveMethod: defaultSaveMethod,
+    unwrapMethod: defaultUnwrapMethod,
+    checkForDuplicateTitle: defaultCheckForDuplicateTitle,
+  };
 
   beforeEach(() => {
     attributes = {
@@ -55,9 +80,10 @@ describe('attributeService', () => {
   });
 
   describe('determining input type', () => {
-    const defaultAttributeService = mockAttributeService<TestAttributes>(defaultTestType);
+    const defaultAttributeService = mockAttributeService<TestAttributes>(defaultTestType, options);
     const customAttributeService = mockAttributeService<TestAttributes, TestByValueInput>(
-      defaultTestType
+      defaultTestType,
+      options
     );
 
     it('can determine input type given default types', () => {
@@ -85,39 +111,32 @@ describe('attributeService', () => {
   });
 
   describe('unwrapping attributes', () => {
-    it('can unwrap all default attributes when given reference type input', async () => {
-      const core = coreMock.createStart();
-      core.savedObjects.client.get = jest.fn().mockResolvedValueOnce({
-        attributes,
+    it('does not throw error when given reference type input with no unwrap method', async () => {
+      const attributeService = mockAttributeService<TestAttributes>(defaultTestType, {
+        saveMethod: defaultSaveMethod,
+        checkForDuplicateTitle: jest.fn(),
       });
-      const attributeService = mockAttributeService<TestAttributes>(
-        defaultTestType,
-        undefined,
-        core
-      );
-      expect(await attributeService.unwrapAttributes(byReferenceInput)).toEqual(attributes);
+      expect(await attributeService.unwrapAttributes(byReferenceInput)).toEqual(byReferenceInput);
     });
 
     it('returns attributes when when given value type input', async () => {
-      const attributeService = mockAttributeService<TestAttributes>(defaultTestType);
+      const attributeService = mockAttributeService<TestAttributes>(defaultTestType, options);
       expect(await attributeService.unwrapAttributes(byValueInput)).toEqual(attributes);
     });
 
     it('runs attributes through a custom unwrap method', async () => {
-      const core = coreMock.createStart();
-      core.savedObjects.client.get = jest.fn().mockResolvedValueOnce({
-        attributes,
-      });
-      const attributeService = mockAttributeService<TestAttributes>(
-        defaultTestType,
-        {
-          customUnwrapMethod: (savedObject) => ({
-            ...savedObject.attributes,
-            testAttr2: { array: [1, 2, 3, 4, 5], testAttr3: 'kibanana' },
-          }),
+      const attributeService = mockAttributeService<TestAttributes>(defaultTestType, {
+        saveMethod: defaultSaveMethod,
+        unwrapMethod: (savedObjectId) => {
+          return new Promise((resolve) => {
+            return resolve({
+              ...attributes,
+              testAttr2: { array: [1, 2, 3, 4, 5], testAttr3: 'kibanana' },
+            });
+          });
         },
-        core
-      );
+        checkForDuplicateTitle: jest.fn(),
+      });
       expect(await attributeService.unwrapAttributes(byReferenceInput)).toEqual({
         ...attributes,
         testAttr2: { array: [1, 2, 3, 4, 5], testAttr3: 'kibanana' },
@@ -127,52 +146,40 @@ describe('attributeService', () => {
 
   describe('wrapping attributes', () => {
     it('returns given attributes when use ref type is false', async () => {
-      const attributeService = mockAttributeService<TestAttributes>(defaultTestType);
+      const attributeService = mockAttributeService<TestAttributes>(defaultTestType, options);
       expect(await attributeService.wrapAttributes(attributes, false)).toEqual({ attributes });
     });
 
-    it('updates existing saved object with new attributes when given id', async () => {
+    it('calls saveMethod with appropriate parameters', async () => {
       const core = coreMock.createStart();
+      const saveMethod = jest.fn();
+      saveMethod.mockReturnValueOnce({});
       const attributeService = mockAttributeService<TestAttributes>(
         defaultTestType,
-        undefined,
+        {
+          saveMethod,
+          unwrapMethod: defaultUnwrapMethod,
+          checkForDuplicateTitle: defaultCheckForDuplicateTitle,
+        },
         core
       );
       expect(await attributeService.wrapAttributes(attributes, true, byReferenceInput)).toEqual(
         byReferenceInput
       );
-      expect(core.savedObjects.client.update).toHaveBeenCalledWith(
-        defaultTestType,
-        '123',
-        attributes
-      );
-    });
-
-    it('creates new saved object with attributes when given no id', async () => {
-      const core = coreMock.createStart();
-      core.savedObjects.client.create = jest.fn().mockResolvedValueOnce({
-        id: '678',
-      });
-      const attributeService = mockAttributeService<TestAttributes>(
-        defaultTestType,
-        undefined,
-        core
-      );
-      expect(await attributeService.wrapAttributes(attributes, true)).toEqual({
-        savedObjectId: '678',
-      });
-      expect(core.savedObjects.client.create).toHaveBeenCalledWith(defaultTestType, attributes);
+      expect(saveMethod).toHaveBeenCalledWith(defaultTestType, attributes, '123');
     });
 
     it('uses custom save method when given an id', async () => {
-      const customSaveMethod = jest.fn().mockReturnValue({ id: '123' });
+      const saveMethod = jest.fn().mockReturnValue({ id: '123' });
       const attributeService = mockAttributeService<TestAttributes>(defaultTestType, {
-        customSaveMethod,
+        saveMethod,
+        unwrapMethod: defaultUnwrapMethod,
+        checkForDuplicateTitle: defaultCheckForDuplicateTitle,
       });
       expect(await attributeService.wrapAttributes(attributes, true, byReferenceInput)).toEqual(
         byReferenceInput
       );
-      expect(customSaveMethod).toHaveBeenCalledWith(
+      expect(saveMethod).toHaveBeenCalledWith(
         defaultTestType,
         attributes,
         byReferenceInput.savedObjectId
@@ -180,14 +187,16 @@ describe('attributeService', () => {
     });
 
     it('uses custom save method given no id', async () => {
-      const customSaveMethod = jest.fn().mockReturnValue({ id: '678' });
+      const saveMethod = jest.fn().mockReturnValue({ id: '678' });
       const attributeService = mockAttributeService<TestAttributes>(defaultTestType, {
-        customSaveMethod,
+        saveMethod,
+        unwrapMethod: defaultUnwrapMethod,
+        checkForDuplicateTitle: defaultCheckForDuplicateTitle,
       });
       expect(await attributeService.wrapAttributes(attributes, true)).toEqual({
         savedObjectId: '678',
       });
-      expect(customSaveMethod).toHaveBeenCalledWith(defaultTestType, attributes, undefined);
+      expect(saveMethod).toHaveBeenCalledWith(defaultTestType, attributes, undefined);
     });
   });
 });

src/plugins/dashboard/public/attribute_service/attribute_service.tsx

@@ -27,22 +27,10 @@ import {
   IEmbeddable,
   Container,
   EmbeddableStart,
-  EmbeddableFactory,
   EmbeddableFactoryNotFoundError,
 } from '../embeddable_plugin';
-import {
-  SavedObjectsClientContract,
-  SimpleSavedObject,
-  I18nStart,
-  NotificationsStart,
-  OverlayStart,
-} from '../../../../core/public';
-import {
-  SavedObjectSaveModal,
-  OnSaveProps,
-  SaveResult,
-  checkForDuplicateTitle,
-} from '../../../saved_objects/public';
+import { I18nStart, NotificationsStart } from '../../../../core/public';
+import { SavedObjectSaveModal, OnSaveProps, SaveResult } from '../../../saved_objects/public';
 
 /**
  * The attribute service is a shared, generic service that embeddables can use to provide the functionality
@@ -53,12 +41,13 @@ import {
 export const ATTRIBUTE_SERVICE_KEY = 'attributes';
 
 export interface AttributeServiceOptions<A extends { title: string }> {
-  customSaveMethod?: (
+  saveMethod: (
     type: string,
     attributes: A,
     savedObjectId?: string
   ) => Promise<{ id?: string } | { error: Error }>;
-  customUnwrapMethod?: (savedObject: SimpleSavedObject<A>) => A;
+  checkForDuplicateTitle: (props: OnSaveProps) => Promise<true>;
+  unwrapMethod?: (savedObjectId: string) => Promise<A>;
 }
 
 export class AttributeService<
@@ -68,38 +57,37 @@ export class AttributeService<
   } = EmbeddableInput & { [ATTRIBUTE_SERVICE_KEY]: SavedObjectAttributes },
   RefType extends SavedObjectEmbeddableInput = SavedObjectEmbeddableInput
 > {
-  private embeddableFactory?: EmbeddableFactory;
-
   constructor(
     private type: string,
     private showSaveModal: (
       saveModal: React.ReactElement,
       I18nContext: I18nStart['Context']
     ) => void,
-    private savedObjectsClient: SavedObjectsClientContract,
-    private overlays: OverlayStart,
     private i18nContext: I18nStart['Context'],
     private toasts: NotificationsStart['toasts'],
-    getEmbeddableFactory?: EmbeddableStart['getEmbeddableFactory'],
-    private options?: AttributeServiceOptions<SavedObjectAttributes>
+    private options: AttributeServiceOptions<SavedObjectAttributes>,
+    getEmbeddableFactory?: EmbeddableStart['getEmbeddableFactory']
   ) {
     if (getEmbeddableFactory) {
       const factory = getEmbeddableFactory(this.type);
       if (!factory) {
         throw new EmbeddableFactoryNotFoundError(this.type);
       }
-      this.embeddableFactory = factory;
     }
   }
 
+  private async defaultUnwrapMethod(input: RefType): Promise<SavedObjectAttributes> {
+    return new Promise<SavedObjectAttributes>((resolve) => {
+      // @ts-ignore
+      return resolve({ ...input });
+    });
+  }
+
   public async unwrapAttributes(input: RefType | ValType): Promise<SavedObjectAttributes> {
     if (this.inputIsRefType(input)) {
-      const savedObject: SimpleSavedObject<SavedObjectAttributes> = await this.savedObjectsClient.get<
-        SavedObjectAttributes
-      >(this.type, input.savedObjectId);
-      return this.options?.customUnwrapMethod
-        ? this.options?.customUnwrapMethod(savedObject)
-        : { ...savedObject.attributes };
+      return this.options.unwrapMethod
+        ? await this.options.unwrapMethod(input.savedObjectId)
+        : await this.defaultUnwrapMethod(input);
     }
     return input[ATTRIBUTE_SERVICE_KEY];
   }
@@ -118,25 +106,11 @@ export class AttributeService<
       return { [ATTRIBUTE_SERVICE_KEY]: newAttributes } as ValType;
     }
     try {
-      if (this.options?.customSaveMethod) {
-        const savedItem = await this.options.customSaveMethod(
-          this.type,
-          newAttributes,
-          savedObjectId
-        );
-        if ('id' in savedItem) {
-          return { ...originalInput, savedObjectId: savedItem.id } as RefType;
-        }
-        return { ...originalInput } as RefType;
+      const savedItem = await this.options.saveMethod(this.type, newAttributes, savedObjectId);
+      if ('id' in savedItem) {
+        return { ...originalInput, savedObjectId: savedItem.id } as RefType;
       }
-
-      if (savedObjectId) {
-        await this.savedObjectsClient.update(this.type, savedObjectId, newAttributes);
-        return { ...originalInput, savedObjectId } as RefType;
-      }
-
-      const savedItem = await this.savedObjectsClient.create(this.type, newAttributes);
-      return { ...originalInput, savedObjectId: savedItem.id } as RefType;
+      return { ...originalInput } as RefType;
     } catch (error) {
       this.toasts.addDanger({
         title: i18n.translate('dashboard.attributeService.saveToLibraryError', {
@@ -181,21 +155,7 @@ export class AttributeService<
     }
     return new Promise<RefType>((resolve, reject) => {
       const onSave = async (props: OnSaveProps): Promise<SaveResult> => {
-        await checkForDuplicateTitle(
-          {
-            title: props.newTitle,
-            copyOnSave: false,
-            lastSavedTitle: '',
-            getEsType: () => this.type,
-            getDisplayName: this.embeddableFactory?.getDisplayName || (() => this.type),
-          },
-          props.isTitleDuplicateConfirmed,
-          props.onTitleDuplicate,
-          {
-            savedObjectsClient: this.savedObjectsClient,
-            overlays: this.overlays,
-          }
-        );
+        await this.options.checkForDuplicateTitle(props);
         try {
           const newAttributes = { ...input[ATTRIBUTE_SERVICE_KEY] };
           newAttributes.title = props.newTitle;

src/plugins/dashboard/public/plugin.tsx

@@ -164,7 +164,7 @@ export interface DashboardStart {
     R extends SavedObjectEmbeddableInput = SavedObjectEmbeddableInput
   >(
     type: string,
-    options?: AttributeServiceOptions<A>
+    options: AttributeServiceOptions<A>
   ) => AttributeService<A, V, R>;
 }
 
@@ -482,12 +482,10 @@ export class DashboardPlugin
         new AttributeService(
           type,
           showSaveModal,
-          core.savedObjects.client,
-          core.overlays,
           core.i18n.Context,
           core.notifications.toasts,
-          embeddable.getEmbeddableFactory,
-          options
+          options,
+          embeddable.getEmbeddableFactory
         ),
     };
   }

src/plugins/visualizations/public/embeddable/visualize_embeddable_factory.tsx

@@ -18,7 +18,7 @@
  */
 
 import { i18n } from '@kbn/i18n';
-import { SavedObjectMetaData } from 'src/plugins/saved_objects/public';
+import { SavedObjectMetaData, OnSaveProps } from 'src/plugins/saved_objects/public';
 import { first } from 'rxjs/operators';
 import { SavedObjectAttributes } from '../../../../core/public';
 import {
@@ -51,6 +51,7 @@ import { StartServicesGetter } from '../../../kibana_utils/public';
 import { VisualizationsStartDeps } from '../plugin';
 import { VISUALIZE_ENABLE_LABS_SETTING } from '../../common/constants';
 import { AttributeService } from '../../../dashboard/public';
+import { checkForDuplicateTitle } from '../../../saved_objects/public';
 
 interface VisualizationAttributes extends SavedObjectAttributes {
   visState: string;
@@ -58,7 +59,7 @@ interface VisualizationAttributes extends SavedObjectAttributes {
 
 export interface VisualizeEmbeddableFactoryDeps {
   start: StartServicesGetter<
-    Pick<VisualizationsStartDeps, 'inspector' | 'embeddable' | 'dashboard'>
+    Pick<VisualizationsStartDeps, 'inspector' | 'embeddable' | 'dashboard' | 'savedObjectsClient'>
   >;
 }
 
@@ -129,7 +130,10 @@ export class VisualizeEmbeddableFactory
           VisualizeSavedObjectAttributes,
           VisualizeByValueInput,
           VisualizeByReferenceInput
-        >(this.type, { customSaveMethod: this.onSave });
+        >(this.type, {
+          saveMethod: this.saveMethod.bind(this),
+          checkForDuplicateTitle: this.checkTitle.bind(this),
+        });
     }
     return this.attributeService!;
   }
@@ -183,7 +187,7 @@ export class VisualizeEmbeddableFactory
     }
   }
 
-  private async onSave(
+  private async saveMethod(
     type: string,
     attributes: VisualizeSavedObjectAttributes
   ): Promise<{ id: string }> {
@@ -225,4 +229,24 @@ export class VisualizeEmbeddableFactory
       throw error;
     }
   }
+
+  public async checkTitle(props: OnSaveProps): Promise<true> {
+    const savedObjectsClient = await this.deps.start().core.savedObjects.client;
+    const overlays = await this.deps.start().core.overlays;
+    return checkForDuplicateTitle(
+      {
+        title: props.newTitle,
+        copyOnSave: false,
+        lastSavedTitle: '',
+        getEsType: () => this.type,
+        getDisplayName: this.getDisplayName || (() => this.type),
+      },
+      props.isTitleDuplicateConfirmed,
+      props.onTitleDuplicate,
+      {
+        savedObjectsClient,
+        overlays,
+      }
+    );
+  }
 }

src/plugins/visualizations/public/mocks.ts

@@ -72,6 +72,7 @@ const createInstance = async () => {
       embeddable: embeddablePluginMock.createStartContract(),
       dashboard: dashboardPluginMock.createStartContract(),
       getAttributeService: jest.fn(),
+      savedObjectsClient: coreMock.createStart().savedObjects.client,
     });
 
   return {

src/plugins/visualizations/public/plugin.ts

@@ -25,6 +25,7 @@ import {
   CoreStart,
   Plugin,
   ApplicationStart,
+  SavedObjectsClientContract,
 } from '../../../core/public';
 import { TypesService, TypesSetup, TypesStart } from './vis_types';
 import {
@@ -112,6 +113,7 @@ export interface VisualizationsStartDeps {
   application: ApplicationStart;
   dashboard: DashboardStart;
   getAttributeService: DashboardStart['getAttributeService'];
+  savedObjectsClient: SavedObjectsClientContract;
 }
 
 /**

x-pack/plugins/lens/public/app_plugin/app.test.tsx

@@ -145,8 +145,9 @@ describe('Lens App', () => {
     >(
       DOC_TYPE,
       {
-        customSaveMethod: jest.fn(),
-        customUnwrapMethod: jest.fn(),
+        saveMethod: jest.fn(),
+        unwrapMethod: jest.fn(),
+        checkForDuplicateTitle: jest.fn(),
       },
       core
     );

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

@@ -27,6 +27,7 @@ import { coreMock, httpServiceMock } from '../../../../../../src/core/public/moc
 import { IBasePath } from '../../../../../../src/core/public';
 import { AttributeService } from '../../../../../../src/plugins/dashboard/public';
 import { LensAttributeService } from '../../lens_attribute_service';
+import { OnSaveProps } from '../../../../../../src/plugins/saved_objects/public/save_modal';
 
 jest.mock('../../../../../../src/plugins/inspector/public/', () => ({
   isAvailable: false,
@@ -44,6 +45,30 @@ const savedVis: Document = {
   title: 'My title',
   visualizationType: '',
 };
+const defaultSaveMethod = (
+  type: string,
+  testAttributes: LensSavedObjectAttributes,
+  savedObjectId?: string
+): Promise<{ id: string }> => {
+  return new Promise(() => {
+    return { id: '123' };
+  });
+};
+const defaultUnwrapMethod = (savedObjectId: string): Promise<LensSavedObjectAttributes> => {
+  return new Promise(() => {
+    return { ...savedVis };
+  });
+};
+const defaultCheckForDuplicateTitle = (props: OnSaveProps): Promise<true> => {
+  return new Promise(() => {
+    return true;
+  });
+};
+const options = {
+  saveMethod: defaultSaveMethod,
+  unwrapMethod: defaultUnwrapMethod,
+  checkForDuplicateTitle: defaultCheckForDuplicateTitle,
+};
 
 const attributeServiceMockFromSavedVis = (document: Document): LensAttributeService => {
   const core = coreMock.createStart();
@@ -51,14 +76,7 @@ const attributeServiceMockFromSavedVis = (document: Document): LensAttributeServ
     LensSavedObjectAttributes,
     LensByValueInput,
     LensByReferenceInput
-  >(
-    'lens',
-    jest.fn(),
-    core.savedObjects.client,
-    core.overlays,
-    core.i18n.Context,
-    core.notifications.toasts
-  );
+  >('lens', jest.fn(), core.i18n.Context, core.notifications.toasts, options);
   service.unwrapAttributes = jest.fn((input: LensByValueInput | LensByReferenceInput) => {
     return Promise.resolve({ ...document } as LensSavedObjectAttributes);
   });

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

@@ -13,6 +13,7 @@ import {
   LensByReferenceInput,
 } from './editor_frame_service/embeddable/embeddable';
 import { SavedObjectIndexStore, DOC_TYPE } from './persistence';
+import { checkForDuplicateTitle, OnSaveProps } from '../../../../src/plugins/saved_objects/public';
 
 export type LensAttributeService = AttributeService<
   LensSavedObjectAttributes,
@@ -30,7 +31,7 @@ export function getLensAttributeService(
     LensByValueInput,
     LensByReferenceInput
   >(DOC_TYPE, {
-    customSaveMethod: async (
+    saveMethod: async (
       type: string,
       attributes: LensSavedObjectAttributes,
       savedObjectId?: string
@@ -42,11 +43,34 @@ export function getLensAttributeService(
       });
       return { id: savedDoc.savedObjectId };
     },
-    customUnwrapMethod: (savedObject) => {
+    unwrapMethod: async (savedObjectId: string): Promise<LensSavedObjectAttributes> => {
+      const savedObject = await core.savedObjects.client.get<LensSavedObjectAttributes>(
+        DOC_TYPE,
+        savedObjectId
+      );
       return {
         ...savedObject.attributes,
         references: savedObject.references,
       };
     },
+    checkForDuplicateTitle: (props: OnSaveProps) => {
+      const savedObjectsClient = core.savedObjects.client;
+      const overlays = core.overlays;
+      return checkForDuplicateTitle(
+        {
+          title: props.newTitle,
+          copyOnSave: false,
+          lastSavedTitle: '',
+          getEsType: () => DOC_TYPE,
+          getDisplayName: () => DOC_TYPE,
+        },
+        props.isTitleDuplicateConfirmed,
+        props.onTitleDuplicate,
+        {
+          savedObjectsClient,
+          overlays,
+        }
+      );
+    },
   });
 }