Using custom save method on attribute service

Author: Maja Grubic

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

@@ -189,5 +189,15 @@ describe('attributeService', () => {
       });
       expect(customSaveMethod).toHaveBeenCalledWith(defaultTestType, attributes, undefined);
     });
+
+    it('uses custom save method when passed through setOptions', async () => {
+      const customSaveMethod = jest.fn().mockReturnValue({ id: '678' });
+      const attributeService = mockAttributeService<TestAttributes>(defaultTestType);
+      attributeService.setOptions({ customSaveMethod });
+      expect(await attributeService.wrapAttributes(attributes, true)).toEqual({
+        savedObjectId: '678',
+      });
+      expect(customSaveMethod).toHaveBeenCalledWith(defaultTestType, attributes, undefined);
+    });
   });
 });

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

@@ -57,7 +57,7 @@ export interface AttributeServiceOptions<A extends { title: string }> {
     type: string,
     attributes: A,
     savedObjectId?: string
-  ) => Promise<{ id: string }>;
+  ) => Promise<{ id?: string } | { error: Error}>;
   customUnwrapMethod?: (savedObject: SimpleSavedObject<A>) => A;
 }
 
@@ -124,7 +124,10 @@ export class AttributeService<
           newAttributes,
           savedObjectId
         );
-        return { ...originalInput, savedObjectId: savedItem.id } as RefType;
+        if ('id' in savedItem) {
+          return { ...originalInput, savedObjectId: savedItem.id } as RefType;
+        }
+        return { ...originalInput } as RefType;
       }
 
       if (savedObjectId) {
@@ -176,7 +179,6 @@ export class AttributeService<
   getInputAsRefType = async (
     input: ValType | RefType,
     saveOptions?: { showSaveModal: boolean; saveModalTitle?: string } | { title: string },
-    customSaveMethod?: (props: OnSaveProps) => Promise<SaveResult>
   ): Promise<RefType> => {
     if (this.inputIsRefType(input)) {
       return input;
@@ -209,11 +211,10 @@ export class AttributeService<
           return { error };
         }
       };
-      const onSaveMethod = customSaveMethod ? customSaveMethod : onSave;
       if (saveOptions && (saveOptions as { showSaveModal: boolean }).showSaveModal) {
         this.showSaveModal(
           <SavedObjectSaveModal
-            onSave={onSaveMethod}
+            onSave={onSave}
             onClose={() => reject()}
             title={get(saveOptions, 'saveModalTitle', input[ATTRIBUTE_SERVICE_KEY].title)}
             showCopyOnSave={false}
@@ -225,4 +226,8 @@ export class AttributeService<
       }
     });
   };
+
+  public setOptions = (options: AttributeServiceOptions<SavedObjectAttributes>) => {
+    this.options = options;
+  };
 }

src/plugins/visualizations/public/embeddable/create_vis_embeddable_from_object.ts

@@ -42,12 +42,12 @@ import { SavedVisualizationsLoader } from '../saved_visualizations';
 export const createVisEmbeddableFromObject = (deps: VisualizeEmbeddableFactoryDeps) => async (
   vis: Vis,
   input: Partial<VisualizeInput> & { id: string },
+  savedVisualizationsLoader?: SavedVisualizationsLoader,
   attributeService?: AttributeService<
     VisualizeSavedObjectAttributes,
     VisualizeByValueInput,
     VisualizeByReferenceInput
   >,
-  savedVisualizationsLoader?: SavedVisualizationsLoader,
   parent?: IContainer
 ): Promise<VisualizeEmbeddable | ErrorEmbeddable | DisabledLabEmbeddable> => {
   const savedVisualizations = getSavedVisualizationsLoader();

src/plugins/visualizations/public/embeddable/visualize_embeddable.ts

@@ -54,7 +54,6 @@ import { SavedObjectAttributes } from '../../../../core/types';
 import { AttributeService } from '../../../dashboard/public';
 import { SavedVisualizationsLoader } from '../saved_visualizations';
 import { VisSavedObject } from '../types';
-import { OnSaveProps, SaveResult } from '../../../saved_objects/public';
 
 const getKeys = <T extends {}>(o: T): Array<keyof T> => Object.keys(o) as Array<keyof T>;
 
@@ -83,7 +82,11 @@ export interface VisualizeOutput extends EmbeddableOutput {
   visTypeName: string;
 }
 
-export type VisualizeSavedObjectAttributes = SavedObjectAttributes & { title: string };
+export type VisualizeSavedObjectAttributes = SavedObjectAttributes & {
+  title: string;
+  vis?: Vis;
+  savedVis?: VisSavedObject;
+};
 export type VisualizeByValueInput = { attributes: VisualizeSavedObjectAttributes } & VisualizeInput;
 export type VisualizeByReferenceInput = SavedObjectEmbeddableInput & VisualizeInput;
 
@@ -426,49 +429,64 @@ export class VisualizeEmbeddable
   };
 
   getInputAsRefType = async (): Promise<VisualizeByReferenceInput> => {
+    const savedVis: VisSavedObject = await this.savedVisualizationsLoader?.get({});
+    if (!savedVis) {
+      throw new Error('No Saved Vis');
+    }
     return new Promise<VisualizeByReferenceInput>((resolve, reject) => {
-      const onSave = async (props: OnSaveProps): Promise<SaveResult> => {
+      const onSave = async (
+        type: string,
+        attributes: VisualizeSavedObjectAttributes
+      ): Promise<{ id: string }> => {
         try {
-          const savedVis: VisSavedObject = await this.savedVisualizationsLoader?.get({});
+          const { title, vis } = attributes;
           const saveOptions = {
             confirmOverwrite: false,
             returnToOrigin: true,
           };
-          savedVis.title = props.newTitle;
-          savedVis.copyOnSave = props.newCopyOnSave || false;
-          savedVis.description = props.newDescription || '';
-          savedVis.searchSourceFields = this.vis?.data.searchSource?.getSerializedFields();
-          const serializedVis = this.vis.serialize();
-          const { title, type, params, data } = serializedVis;
+          savedVis.title = title;
+          savedVis.copyOnSave = false;
+          savedVis.description = '';
+          savedVis.searchSourceFields = vis?.data.searchSource?.getSerializedFields();
+          const serializedVis = ((vis as unknown) as Vis).serialize();
+          const { params, data } = serializedVis;
           savedVis.visState = {
             title,
-            type,
+            type: serializedVis.type,
             params,
             aggs: data.aggs,
           };
-          savedVis.uiStateJSON = this.vis.uiState.toString();
+          if (vis) {
+            savedVis.uiStateJSON = vis.uiState.toString();
+          }
           const id = await savedVis.save(saveOptions);
           resolve({ savedObjectId: id, id: this.id });
           return { id };
         } catch (error) {
           reject(error);
-          return { error };
+          return { id: '' };
         }
       };
       const saveModalTitle = this.getTitle()
         ? this.getTitle()
         : i18n.translate('visualize.embeddable.placeholderTitle', {
             defaultMessage: 'Placeholder Title',
           });
+      // @ts-ignore
+      const attributes: VisualizeSavedObjectAttributes = {
+        savedVis,
+        vis: this.vis,
+        title: this.vis.title,
+      };
+      this.attributeService.setOptions({
+        customSaveMethod: onSave,
+      });
       return this.attributeService.getInputAsRefType(
         {
           id: this.id,
-          attributes: {
-            title: this.vis.title,
-          },
+          attributes,
         },
-        { showSaveModal: true, saveModalTitle },
-        onSave
+        { showSaveModal: true, saveModalTitle }
       );
     });
   };

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

@@ -149,8 +149,8 @@ export class VisualizeEmbeddableFactory
       return createVisEmbeddableFromObject(this.deps)(
         vis,
         input,
-        await this.getAttributeService(),
         savedVisualizations,
+        await this.getAttributeService(),
         parent
       );
     } catch (e) {
@@ -170,8 +170,8 @@ export class VisualizeEmbeddableFactory
       return createVisEmbeddableFromObject(this.deps)(
         vis,
         input,
-        await this.getAttributeService(),
         savedVisualizations,
+        await this.getAttributeService(),
         parent
       );
     } else {