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

@@ -6,13 +6,19 @@
 import { ExpressionsSetup } from '../../expressions/public';
 import { PluginInitializerContext, CoreSetup, CoreStart, Plugin } from '../../../core/public';
 import { DataPublicPluginSetup, DataPublicPluginStart } from '../../data/public';
+<<<<<<< HEAD
 import { visLayers } from './expressions';
+=======
+import { setSavedAugmentVisLoader } from './services';
+import { createSavedAugmentVisLoader, SavedAugmentVisLoader } from './saved_augment_vis';
+>>>>>>> c3e46d9dc9 (Add augment-vis saved obj)
 
 // 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;
@@ -37,7 +43,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,66 @@
+/*
+ * 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',
+      visId: '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', ''),
+          visId: get(opts, 'visId', ''),
+          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.test.ts

@@ -0,0 +1,41 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { VisLayerExpressionFn } from '../types';
+import {
+  createSavedAugmentVisLoader,
+  SavedObjectOpenSearchDashboardsServicesWithAugmentVis,
+} from './saved_augment_vis';
+import { generateAugmentVisSavedObject, getMockAugmentVisSavedObjectClient } from './utils';
+
+describe('SavedObjectLoaderAugmentVis', () => {
+  const fn = {
+    // TODO: VisLayerTypes will resolve after rebasing with earlier PR
+    type: VisLayerTypes.PointInTimeEventsLayer,
+    name: 'test-fn',
+    args: {
+      testArg: 'test-value',
+    },
+  } as VisLayerExpressionFn;
+  const obj1 = generateAugmentVisSavedObject('test-id-1', fn);
+  const obj2 = generateAugmentVisSavedObject('test-id-2', fn);
+
+  it('findAll returns single saved obj', async () => {
+    const loader = createSavedAugmentVisLoader({
+      savedObjectsClient: getMockAugmentVisSavedObjectClient([obj1]),
+    } as SavedObjectOpenSearchDashboardsServicesWithAugmentVis);
+    const resp = await loader.findAll();
+    expect(resp.hits.length === 1);
+  });
+
+  // TODO: once done rebasing after VisLayer PR, can finish creating test cases here.
+  // right now they are failing since there is missing imports
+
+  // add test for empty response
+  // add test for multi obj response
+  // add test for invalid VisLayerType
+  // add test for missing reference
+  // add test for missing visLayerExpressionFn
+});

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

@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { get, isEmpty } from 'lodash';
+import {
+  SavedObjectLoader,
+  SavedObjectOpenSearchDashboardsServices,
+} from '../../../saved_objects/public';
+import { createSavedAugmentVisClass } from './_saved_augment_vis';
+
+// 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 {
+    mapHitSource = (source: Record<string, any>, id: string) => {
+      source.id = id;
+      source.visId = get(source, 'visReference.id', '');
+
+      if (isEmpty(source.visReference)) {
+        source.error = 'visReference is missing in augment-vis saved object';
+      }
+      if (isEmpty(source.visLayerExpressionFn)) {
+        source.error = 'visLayerExpressionFn is missing in augment-vis saved object';
+      }
+      // TODO: will resolve after rebasing with earlier PR
+      if (!(get(source, 'visLayerExpressionFn.type', '') in VisLayerTypes)) {
+        source.error = 'Unknown VisLayer expression function type';
+      }
+
+      delete source.visReference;
+      delete source.visName;
+      return source;
+    };
+
+    /**
+     * Updates hit.attributes to contain an id related to the referenced visualization
+     * (visId) 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 vis reference per saved object.
+      // If we change to multiple, we will need to dynamically handle that
+      const visReference = hit.references[0];
+      return this.mapHitSource({ ...hit.attributes, visReference }, 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,110 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { extractReferences, injectReferences } from './saved_augment_vis_references';
+import { AugmentVisSavedObject } from '../types';
+
+describe('extractReferences()', () => {
+  test('extracts nothing if visId is null', () => {
+    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 visId', () => {
+    const doc = {
+      id: '1',
+      attributes: {
+        foo: true,
+        visId: 'test-id',
+      },
+      references: [],
+    };
+    const updatedDoc = extractReferences(doc);
+    expect(updatedDoc).toMatchInlineSnapshot(`
+      Object {
+        "attributes": Object {
+          "foo": true,
+          "visName": "visualization_0",
+        },
+        "references": Array [
+          Object {
+            "id": "test-id",
+            "name": "visualization_0",
+            "type": "visualization",
+          },
+        ],
+      }
+    `);
+  });
+});
+
+describe('injectReferences()', () => {
+  test('injects nothing when visName 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',
+      visName: 'visualization_0',
+    } as unknown) as AugmentVisSavedObject;
+    const references = [
+      {
+        name: 'visualization_0',
+        type: 'visualization',
+        id: 'test-id',
+      },
+    ];
+    injectReferences(context, references);
+    expect(context).toMatchInlineSnapshot(`
+      Object {
+        "id": "1",
+        "pluginResourceId": "test-resource-id",
+        "visId": "test-id",
+        "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',
+      visName: 'visualization_0',
+    } as unknown) as AugmentVisSavedObject;
+    expect(() => injectReferences(context, [])).toThrowErrorMatchingInlineSnapshot(
+      `"Could not find visualization reference \\"visualization_0\\""`
+    );
+  });
+});