From f5d07b7c79784df0122c8f1144c0408d0e13aead Mon Sep 17 00:00:00 2001
From: Amit Galitzky <amgalitz@amazon.com>
Date: Wed, 15 Mar 2023 20:19:11 -0700
Subject: [PATCH] setting augmented viz loader and importing as plugin

---
 src/plugins/visualizations/opensearch_dashboards.json | 2 +-
 src/plugins/visualizations/public/plugin.ts           | 6 +++++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/plugins/visualizations/opensearch_dashboards.json b/src/plugins/visualizations/opensearch_dashboards.json
index b7c5e4ab9b4e..b9871161accc 100644
--- a/src/plugins/visualizations/opensearch_dashboards.json
+++ b/src/plugins/visualizations/opensearch_dashboards.json
@@ -3,7 +3,7 @@
   "version": "opensearchDashboards",
   "server": true,
   "ui": true,
-  "requiredPlugins": ["data", "expressions", "uiActions", "embeddable", "inspector", "dashboard"],
+  "requiredPlugins": ["data", "expressions", "uiActions", "embeddable", "inspector", "dashboard", "visAugmenter"],
   "optionalPlugins": ["usageCollection"],
   "requiredBundles": ["opensearchDashboardsUtils", "discover", "savedObjects", "visAugmenter"]
 }
diff --git a/src/plugins/visualizations/public/plugin.ts b/src/plugins/visualizations/public/plugin.ts
index 682e678ed584..64e7dfe1d654 100644
--- a/src/plugins/visualizations/public/plugin.ts
+++ b/src/plugins/visualizations/public/plugin.ts
@@ -30,6 +30,7 @@
 
 import './index.scss';
 
+import { VisAugmenterStart } from 'src/plugins/vis_augmenter/public';
 import {
   PluginInitializerContext,
   CoreSetup,
@@ -60,6 +61,7 @@ import {
   setOverlays,
   setSavedSearchLoader,
   setEmbeddable,
+  setSavedAugmentVisLoader,
 } from './services';
 import {
   VISUALIZE_EMBEDDABLE_TYPE,
@@ -128,6 +130,7 @@ export interface VisualizationsStartDeps {
   dashboard: DashboardStart;
   getAttributeService: DashboardStart['getAttributeService'];
   savedObjectsClient: SavedObjectsClientContract;
+  visAugmenter: VisAugmenterStart;
 }
 
 /**
@@ -174,7 +177,7 @@ export class VisualizationsPlugin
 
   public start(
     core: CoreStart,
-    { data, expressions, uiActions, embeddable, dashboard }: VisualizationsStartDeps
+    { data, expressions, uiActions, embeddable, dashboard, visAugmenter }: VisualizationsStartDeps
   ): VisualizationsStart {
     const types = this.types.start();
     setI18n(core.i18n);
@@ -193,6 +196,7 @@ export class VisualizationsPlugin
     setAggs(data.search.aggs);
     setOverlays(core.overlays);
     setChrome(core.chrome);
+    setSavedAugmentVisLoader(visAugmenter.savedAugmentVisLoader);
     const savedVisualizationsLoader = createSavedVisLoader({
       savedObjectsClient: core.savedObjects.client,
       indexPatterns: data.indexPatterns,