Add augment-vis saved obj

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>

src/plugins/vis_augmenter/public/index.ts

@@ -10,3 +10,7 @@ export function plugin(initializerContext: PluginInitializerContext) {
   return new VisAugmenterPlugin(initializerContext);
 }
 export { VisAugmenterSetup, VisAugmenterStart };
+
+export * from './saved_augment_vis';
+
+export { ISavedAugmentVis, VisLayerExpressionFn, AugmentVisSavedObject } from './types';

src/plugins/vis_augmenter/public/plugin.ts

@@ -5,12 +5,15 @@
 
 import { PluginInitializerContext, CoreSetup, CoreStart, Plugin } from '../../../core/public';
 import { DataPublicPluginSetup, DataPublicPluginStart } from '../../data/public';
+import { setSavedAugmentVisLoader } from './services';
+import { createSavedAugmentVisLoader, SavedAugmentVisLoader } from './saved_augment_vis';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
 export interface VisAugmenterSetup {}
 
-// eslint-disable-next-line @typescript-eslint/no-empty-interface
-export interface VisAugmenterStart {}
+export interface VisAugmenterStart {
+  savedAugmentVisLoader: SavedAugmentVisLoader;
+}
 
 export interface VisAugmenterSetupDeps {
   data: DataPublicPluginSetup;
@@ -33,7 +36,15 @@ export class VisAugmenterPlugin
   }
 
   public start(core: CoreStart, { data }: VisAugmenterStartDeps): VisAugmenterStart {
-    return {};
+    const savedAugmentVisLoader = createSavedAugmentVisLoader({
+      savedObjectsClient: core.savedObjects.client,
+      indexPatterns: data.indexPatterns,
+      search: data.search,
+      chrome: core.chrome,
+      overlays: core.overlays,
+    });
+    setSavedAugmentVisLoader(savedAugmentVisLoader);
+    return { savedAugmentVisLoader };
   }
 
   public stop() {}

src/plugins/vis_augmenter/public/saved_augment_vis/_saved_augment_vis.ts

@@ -0,0 +1,68 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/**
+ * @name SavedAugmentVis
+ *
+ * @extends SavedObject.
+ */
+import { get } from 'lodash';
+import {
+  createSavedObjectClass,
+  SavedObject,
+  SavedObjectOpenSearchDashboardsServices,
+} from '../../../saved_objects/public';
+import { IIndexPattern } from '../../../data/public';
+import { extractReferences, injectReferences } from './saved_augment_vis_references';
+
+export function createSavedAugmentVisClass(services: SavedObjectOpenSearchDashboardsServices) {
+  const SavedObjectClass = createSavedObjectClass(services);
+
+  class SavedAugmentVis extends SavedObjectClass {
+    public static type: string = 'augment-vis';
+    public static mapping: Record<string, string> = {
+      description: 'text',
+      pluginResourceId: 'text',
+      savedObjectType: 'keyword',
+      savedObjectId: 'keyword',
+      visLayerExpressionFn: 'object',
+      version: 'integer',
+    };
+
+    constructor(opts: Record<string, unknown> | string = {}) {
+      if (typeof opts !== 'object') {
+        opts = { id: opts };
+      }
+      super({
+        type: SavedAugmentVis.type,
+        mapping: SavedAugmentVis.mapping,
+        extractReferences,
+        injectReferences,
+        id: (opts.id as string) || '',
+        indexPattern: opts.indexPattern as IIndexPattern,
+        defaults: {
+          description: get(opts, 'description', ''),
+          pluginResourceId: get(opts, 'pluginResourceId', ''),
+          savedObjectType: get(opts, 'savedObjectType', ''),
+          savedObjectId: get(opts, 'savedObjectId', ''),
+          visLayerExpressionFn: get(opts, 'visLayerExpressionFn', {}),
+          version: 1,
+        },
+      });
+      // TODO: determine if this saved obj should be visible in saved_obj_management plugin.
+      // if not, we can set showInRecentlyAccessed to false and not persist any edit URL
+      // probably set to false since this saved obj should be hidden by default
+      this.showInRecentlyAccessed = false;
+
+      // we probably don't need this below field. we aren't going to need a full path
+      // since we aren't going to allow editing by default
+      //   this.getFullPath = () => {
+      //     return `/app/visualize#/edit/${this.id}`;
+      //   };
+    }
+  }
+
+  return SavedAugmentVis as new (opts: Record<string, unknown> | string) => SavedObject;
+}

src/plugins/vis_augmenter/public/saved_augment_vis/index.ts

@@ -0,0 +1,7 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export * from './saved_augment_vis';
+export * from './utils';

src/plugins/vis_augmenter/public/saved_augment_vis/saved_augment_vis.ts

@@ -0,0 +1,58 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { get } from 'lodash';
+import {
+  SavedObjectLoader,
+  SavedObjectOpenSearchDashboardsServices,
+} from '../../../saved_objects/public';
+import { createSavedAugmentVisClass } from './_saved_augment_vis';
+
+// TODO: determine if there should be added enforcement on the augment fn types, and reject
+// saved objs that don't have a valid one. Currently, visualizations does something similar
+// in the saved vis loader.
+
+// eslint-disable-next-line @typescript-eslint/no-empty-interface
+export interface SavedObjectOpenSearchDashboardsServicesWithAugmentVis
+  extends SavedObjectOpenSearchDashboardsServices {}
+export type SavedAugmentVisLoader = ReturnType<typeof createSavedAugmentVisLoader>;
+export function createSavedAugmentVisLoader(
+  services: SavedObjectOpenSearchDashboardsServicesWithAugmentVis
+) {
+  const { savedObjectsClient } = services;
+
+  class SavedObjectLoaderAugmentVis extends SavedObjectLoader {
+    // TODO: below fn is extracting the id & type from the source. Is this needed
+    // when we already have inject/extract reference fns?
+    // need to confirm if this codepath is different than the ones using the reference fns
+    mapHitSource = (source: Record<string, any>, id: string) => {
+      source.id = id;
+      source.savedObjectId = get(source, 'savedObjectReference.id', '');
+      source.savedObjectType = get(source, 'savedObjectReference.type', '');
+      delete source.savedObjectReference;
+      delete source.savedObjectName;
+      return source;
+    };
+
+    /**
+     * Updates hit.attributes to contain an id and fields related to the referenced saved object
+     * (savedObjectId, savedObjectType) and returns the updated attributes object.
+     * @param hit
+     * @returns {hit.attributes} The modified hit.attributes object, with an id and url field.
+     */
+    mapSavedObjectApiHits(hit: {
+      references: any[];
+      attributes: Record<string, unknown>;
+      id: string;
+    }) {
+      // For now we are assuming only one reference per saved object.
+      // If we change to multiple, we will need to dynamically handle that
+      const savedObjectReference = hit.references[0];
+      return this.mapHitSource({ ...hit.attributes, savedObjectReference }, hit.id);
+    }
+  }
+  const SavedAugmentVis = createSavedAugmentVisClass(services);
+  return new SavedObjectLoaderAugmentVis(SavedAugmentVis, savedObjectsClient) as SavedObjectLoader;
+}

src/plugins/vis_augmenter/public/saved_augment_vis/saved_augment_vis_references.test.ts

@@ -0,0 +1,154 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { extractReferences, injectReferences } from './saved_augment_vis_references';
+import { AugmentVisSavedObject } from '../../common';
+
+describe('extractReferences()', () => {
+  test('extracts nothing if savedObjectType is empty', () => {
+    const doc = {
+      id: '1',
+      attributes: {
+        foo: true,
+        savedObjectId: 'test-id',
+      },
+      references: [],
+    };
+    const updatedDoc = extractReferences(doc);
+    expect(updatedDoc).toMatchInlineSnapshot(`
+      Object {
+        "attributes": Object {
+          "foo": true,
+          "savedObjectId": "test-id",
+        },
+        "references": Array [],
+      }
+    `);
+  });
+
+  test('extracts nothing if savedObjectId is empty', () => {
+    const doc = {
+      id: '1',
+      attributes: {
+        foo: true,
+        savedObjectType: 'test-type',
+      },
+      references: [],
+    };
+    const updatedDoc = extractReferences(doc);
+    expect(updatedDoc).toMatchInlineSnapshot(`
+      Object {
+        "attributes": Object {
+          "foo": true,
+          "savedObjectType": "test-type",
+        },
+        "references": Array [],
+      }
+    `);
+  });
+
+  test('extracts nothing if savedObjectType & savedObjectId are empty', () => {
+    const doc = {
+      id: '1',
+      attributes: {
+        foo: true,
+      },
+      references: [],
+    };
+    const updatedDoc = extractReferences(doc);
+    expect(updatedDoc).toMatchInlineSnapshot(`
+      Object {
+        "attributes": Object {
+          "foo": true,
+        },
+        "references": Array [],
+      }
+    `);
+  });
+
+  test('extracts references from savedObjectType & savedObjectId', () => {
+    const doc = {
+      id: '1',
+      attributes: {
+        foo: true,
+        savedObjectType: 'test-type',
+        savedObjectId: 'test-id',
+      },
+      references: [],
+    };
+    const updatedDoc = extractReferences(doc);
+    expect(updatedDoc).toMatchInlineSnapshot(`
+      Object {
+        "attributes": Object {
+          "foo": true,
+          "savedObjectName": "saved_object_0",
+        },
+        "references": Array [
+          Object {
+            "id": "test-id",
+            "name": "saved_object_0",
+            "type": "test-type",
+          },
+        ],
+      }
+    `);
+  });
+});
+
+describe('injectReferences()', () => {
+  test('injects nothing when savedObjectName is null', () => {
+    const context = ({
+      id: '1',
+      pluginResourceId: 'test-resource-id',
+      visLayerExpressionFn: 'test-fn',
+    } as unknown) as AugmentVisSavedObject;
+    injectReferences(context, []);
+    expect(context).toMatchInlineSnapshot(`
+      Object {
+        "id": "1",
+        "pluginResourceId": "test-resource-id",
+        "visLayerExpressionFn": "test-fn",
+      }
+    `);
+  });
+
+  test('injects references into context', () => {
+    const context = ({
+      id: '1',
+      pluginResourceId: 'test-resource-id',
+      visLayerExpressionFn: 'test-fn',
+      savedObjectName: 'saved_object_0',
+    } as unknown) as AugmentVisSavedObject;
+    const references = [
+      {
+        name: 'saved_object_0',
+        type: 'visualization',
+        id: 'test-id',
+      },
+    ];
+    injectReferences(context, references);
+    expect(context).toMatchInlineSnapshot(`
+      Object {
+        "id": "1",
+        "pluginResourceId": "test-resource-id",
+        "savedObjectId": "test-id",
+        "savedObjectType": "visualization",
+        "visLayerExpressionFn": "test-fn",
+      }
+    `);
+  });
+
+  test(`fails when it can't find the saved object reference in the array`, () => {
+    const context = ({
+      id: '1',
+      pluginResourceId: 'test-resource-id',
+      visLayerExpressionFn: 'test-fn',
+      savedObjectName: 'saved_object_0',
+    } as unknown) as AugmentVisSavedObject;
+    expect(() => injectReferences(context, [])).toThrowErrorMatchingInlineSnapshot(
+      `"Could not find saved object reference \\"saved_object_0\\""`
+    );
+  });
+});