From 26b241d48c70b7d6133fbd782cb8dc29ede6b888 Mon Sep 17 00:00:00 2001
From: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>
Date: Wed, 24 Feb 2021 18:32:07 -0500
Subject: [PATCH] [Dashboard] Export appropriate references from byValue panels
 (#91567) (#92713)

* Adds references from byValue panels when saving dashboard

* Remove extra spaces

* Rework a type check

* Fix type check

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

Co-authored-by: Corey Robertson <corey.robertson@elastic.co>
---
 .../visualize_embeddable_factory.tsx          | 40 +++++++++++++++++++
 .../embeddable/embeddable_factory.ts          | 14 ++++++-
 .../embeddable/map_embeddable_factory.ts      | 17 +++++++-
 3 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/src/plugins/visualizations/public/embeddable/visualize_embeddable_factory.tsx b/src/plugins/visualizations/public/embeddable/visualize_embeddable_factory.tsx
index a7a5b8626914a..349e024f31c31 100644
--- a/src/plugins/visualizations/public/embeddable/visualize_embeddable_factory.tsx
+++ b/src/plugins/visualizations/public/embeddable/visualize_embeddable_factory.tsx
@@ -9,7 +9,9 @@
 import { i18n } from '@kbn/i18n';
 import { SavedObjectMetaData, OnSaveProps } from 'src/plugins/saved_objects/public';
 import { first } from 'rxjs/operators';
+import { EmbeddableStateWithType } from 'src/plugins/embeddable/common';
 import { SavedObjectAttributes } from '../../../../core/public';
+import { extractSearchSourceReferences } from '../../../data/public';
 import {
   EmbeddableFactoryDefinition,
   EmbeddableOutput,
@@ -236,4 +238,42 @@ export class VisualizeEmbeddableFactory
       }
     );
   }
+
+  public extract(_state: EmbeddableStateWithType) {
+    const state = (_state as unknown) as VisualizeInput;
+    const references = [];
+
+    if (state.savedVis?.data.searchSource) {
+      const [, searchSourceReferences] = extractSearchSourceReferences(
+        state.savedVis.data.searchSource
+      );
+
+      references.push(...searchSourceReferences);
+    }
+
+    if (state.savedVis?.data.savedSearchId) {
+      references.push({
+        name: 'search_0',
+        type: 'search',
+        id: String(state.savedVis.data.savedSearchId),
+      });
+    }
+
+    if (state.savedVis?.params.controls) {
+      const controls = state.savedVis.params.controls;
+      controls.forEach((control: Record<string, string>, i: number) => {
+        if (!control.indexPattern) {
+          return;
+        }
+        control.indexPatternRefName = `control_${i}_index_pattern`;
+        references.push({
+          name: control.indexPatternRefName,
+          type: 'index-pattern',
+          id: control.indexPattern,
+        });
+      });
+    }
+
+    return { state: _state, references };
+  }
 }
diff --git a/x-pack/plugins/lens/public/editor_frame_service/embeddable/embeddable_factory.ts b/x-pack/plugins/lens/public/editor_frame_service/embeddable/embeddable_factory.ts
index 4c40282012d6d..a676b7283671c 100644
--- a/x-pack/plugins/lens/public/editor_frame_service/embeddable/embeddable_factory.ts
+++ b/x-pack/plugins/lens/public/editor_frame_service/embeddable/embeddable_factory.ts
@@ -5,10 +5,11 @@
  * 2.0.
  */
 
-import { Capabilities, HttpSetup } from 'kibana/public';
+import { Capabilities, HttpSetup, SavedObjectReference } from 'kibana/public';
 import { i18n } from '@kbn/i18n';
 import { RecursiveReadonly } from '@kbn/utility-types';
 import { Ast } from '@kbn/interpreter/target/common';
+import { EmbeddableStateWithType } from 'src/plugins/embeddable/common';
 import {
   IndexPatternsContract,
   TimefilterContract,
@@ -105,4 +106,15 @@ export class EmbeddableFactory implements EmbeddableFactoryDefinition {
       parent
     );
   }
+
+  extract(state: EmbeddableStateWithType) {
+    let references: SavedObjectReference[] = [];
+    const typedState = (state as unknown) as LensEmbeddableInput;
+
+    if ('attributes' in typedState && typedState.attributes !== undefined) {
+      references = typedState.attributes.references;
+    }
+
+    return { state, references };
+  }
 }
diff --git a/x-pack/plugins/maps/public/embeddable/map_embeddable_factory.ts b/x-pack/plugins/maps/public/embeddable/map_embeddable_factory.ts
index b039076305498..7e15bfa9a340e 100644
--- a/x-pack/plugins/maps/public/embeddable/map_embeddable_factory.ts
+++ b/x-pack/plugins/maps/public/embeddable/map_embeddable_factory.ts
@@ -6,6 +6,7 @@
  */
 
 import { i18n } from '@kbn/i18n';
+import { EmbeddableStateWithType } from 'src/plugins/embeddable/common';
 import {
   EmbeddableFactoryDefinition,
   IContainer,
@@ -13,8 +14,10 @@ import {
 import '../index.scss';
 import { MAP_SAVED_OBJECT_TYPE, APP_ICON } from '../../common/constants';
 import { getMapEmbeddableDisplayName } from '../../common/i18n_getters';
-import { MapByReferenceInput, MapEmbeddableInput } from './types';
+import { MapByReferenceInput, MapEmbeddableInput, MapByValueInput } from './types';
 import { lazyLoadMapModules } from '../lazy_load_bundle';
+// @ts-expect-error
+import { extractReferences } from '../../common/migrations/references';
 
 export class MapEmbeddableFactory implements EmbeddableFactoryDefinition {
   type = MAP_SAVED_OBJECT_TYPE;
@@ -61,4 +64,16 @@ export class MapEmbeddableFactory implements EmbeddableFactoryDefinition {
       parent
     );
   };
+
+  extract(state: EmbeddableStateWithType) {
+    const maybeMapByValueInput = state as EmbeddableStateWithType | MapByValueInput;
+
+    if ((maybeMapByValueInput as MapByValueInput).attributes !== undefined) {
+      const { references } = extractReferences(maybeMapByValueInput);
+
+      return { state, references };
+    }
+
+    return { state, references: [] };
+  }
 }