[MDS] Install Vega sample data (opensearch-project#6218)

* Add data_source_name to sample data install Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Add unit test coverage Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Resolve CHANGELOG.md conflicts Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Refactor logic Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Remove newline Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Remove newline Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Add comments for spacing Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> * Update jsdoc Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> --------- Signed-off-by: Huy Nguyen <73027756+huyaboo@users.noreply.github.com> (cherry picked from commit 8810f08)
huyaboo · Mar 29, 2024 · 71c1dc9 · 71c1dc9
1 parent a287203
commit 71c1dc9
Show file tree

Hide file tree

Showing 7 changed files with 284 additions and 1 deletion.
diff --git a/src/core/server/index.ts b/src/core/server/index.ts
@@ -319,6 +319,8 @@ export {
   exportSavedObjectsToStream,
   importSavedObjectsFromStream,
   resolveSavedObjectsImportErrors,
+  updateDataSourceNameInVegaSpec,
+  extractVegaSpecFromSavedObject,
 } from './saved_objects';
 
 export {

diff --git a/src/core/server/saved_objects/import/index.ts b/src/core/server/saved_objects/import/index.ts
@@ -43,3 +43,4 @@ export {
   SavedObjectsResolveImportErrorsOptions,
   SavedObjectsImportRetry,
 } from './types';
+export { updateDataSourceNameInVegaSpec, extractVegaSpecFromSavedObject } from './utils';
diff --git a/src/core/server/saved_objects/import/utils.ts b/src/core/server/saved_objects/import/utils.ts
@@ -0,0 +1,113 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { parse, stringify } from 'hjson';
+import { SavedObject, SavedObjectsClientContract } from '../types';
+
+/**
+ * Given a Vega spec, the new datasource (by name), and spacing, update the Vega spec to add the new datasource name to each local cluster query
+ *
+ * @param {string} spec - the stringified Vega spec (HJSON or JSON)
+ * @param {string} newDataSourceName - the datasource name to append
+ * @param {number} [spacing=2] - how large the indenting should be after updating the spec (should be set to > 0 for a readable spec)
+ */
+export interface UpdateDataSourceNameInVegaSpecProps {
+  spec: string;
+  newDataSourceName: string;
+  spacing?: number;
+}
+
+export const updateDataSourceNameInVegaSpec = (
+  props: UpdateDataSourceNameInVegaSpecProps
+): string => {
+  const { spec, spacing } = props;
+  const stringifiedSpacing = spacing || 2;
+
+  let parsedSpec = parseJSONSpec(spec);
+  const isJSONString = !!parsedSpec;
+  if (!parsedSpec) {
+    parsedSpec = parse(spec, { keepWsc: true });
+  }
+
+  const dataField = parsedSpec.data;
+
+  if (dataField instanceof Array) {
+    parsedSpec.data = dataField.map((dataObject) => {
+      return updateDataSourceNameForDataObject(dataObject, props);
+    });
+  } else if (dataField instanceof Object) {
+    parsedSpec.data = updateDataSourceNameForDataObject(dataField, props);
+  } else {
+    throw new Error(`"data" field should be an object or an array of objects`);
+  }
+
+  return isJSONString
+    ? JSON.stringify(parsedSpec)
+    : stringify(parsedSpec, {
+        bracesSameLine: true,
+        keepWsc: true,
+        space: stringifiedSpacing,
+      });
+};
+
+export const getDataSourceTitleFromId = async (
+  dataSourceId: string,
+  savedObjectsClient: SavedObjectsClientContract
+) => {
+  return await savedObjectsClient.get('data-source', dataSourceId).then((response) => {
+    // @ts-expect-error
+    return response?.attributes?.title ?? undefined;
+  });
+};
+
+export const extractVegaSpecFromSavedObject = (savedObject: SavedObject) => {
+  if (isVegaVisualization(savedObject)) {
+    // @ts-expect-error
+    const visStateObject = JSON.parse(savedObject.attributes?.visState);
+    return visStateObject.params.spec;
+  }
+
+  return undefined;
+};
+
+const isVegaVisualization = (savedObject: SavedObject) => {
+  // @ts-expect-error
+  const visState = savedObject.attributes?.visState;
+  if (!!visState) {
+    const visStateObject = JSON.parse(visState);
+    return !!visStateObject.type && visStateObject.type === 'vega';
+  }
+  return false;
+};
+
+const updateDataSourceNameForDataObject = (
+  dataObject: any,
+  props: UpdateDataSourceNameInVegaSpecProps
+) => {
+  const { newDataSourceName } = props;
+  if (
+    dataObject.hasOwnProperty('url') &&
+    dataObject.url.hasOwnProperty('index') &&
+    !dataObject.url.hasOwnProperty('data_source_name')
+  ) {
+    dataObject.url.data_source_name = newDataSourceName;
+  }
+
+  return dataObject;
+};
+
+const parseJSONSpec = (spec: string) => {
+  try {
+    const jsonSpec = JSON.parse(spec);
+
+    if (jsonSpec && typeof jsonSpec === 'object') {
+      return jsonSpec;
+    }
+  } catch (e) {
+    return undefined;
+  }
+
+  return undefined;
+};
diff --git a/src/plugins/home/server/services/sample_data/data_sets/logs/saved_objects.ts b/src/plugins/home/server/services/sample_data/data_sets/logs/saved_objects.ts
diff --git a/src/plugins/home/server/services/sample_data/data_sets/test_utils/visualization_objects.json b/src/plugins/home/server/services/sample_data/data_sets/test_utils/visualization_objects.json
@@ -0,0 +1,73 @@
+{
+    "saved_objects": [
+        {
+            "id": "935afa20-e0cd-11e7-9d07-1398ccfcefa3",
+            "type": "visualization",
+            "updated_at": "2018-08-29T13:22:17.617Z",
+            "version": "1",
+            "migrationVersion": {},
+            "attributes": {
+                "title": "[Logs] Heatmap",
+                "visState": "{\"title\":\"[Logs] Heatmap\",\"type\":\"heatmap\",\"params\":{\"type\":\"heatmap\",\"addTooltip\":true,\"addLegend\":true,\"enableHover\":true,\"legendPosition\":\"right\",\"times\":[],\"colorsNumber\":10,\"colorSchema\":\"Reds\",\"setColorRange\":false,\"colorsRange\":[],\"invertColors\":false,\"percentageMode\":false,\"valueAxes\":[{\"show\":false,\"id\":\"ValueAxis-1\",\"type\":\"value\",\"scale\":{\"type\":\"linear\",\"defaultYExtents\":false},\"labels\":{\"show\":false,\"rotate\":0,\"color\":\"#555\",\"overwriteColor\":false}}]},\"aggs\":[{\"id\":\"1\",\"enabled\":true,\"type\":\"cardinality\",\"schema\":\"metric\",\"params\":{\"field\":\"clientip\"}},{\"id\":\"3\",\"enabled\":true,\"type\":\"terms\",\"schema\":\"group\",\"params\":{\"field\":\"geo.src\",\"size\":5,\"order\":\"desc\",\"orderBy\":\"1\",\"otherBucket\":false,\"otherBucketLabel\":\"Other\",\"missingBucket\":false,\"missingBucketLabel\":\"Missing\",\"customLabel\":\"Country Source\"}},{\"id\":\"2\",\"enabled\":true,\"type\":\"terms\",\"schema\":\"segment\",\"params\":{\"field\":\"hour_of_day\",\"size\":25,\"order\":\"asc\",\"orderBy\":\"_key\",\"otherBucket\":false,\"otherBucketLabel\":\"Other\",\"missingBucket\":false,\"missingBucketLabel\":\"Missing\",\"customLabel\":\"Hour of Day\"}}]}",
+                "uiStateJSON": "{}",
+                "description": "",
+                "version": 1,
+                "kibanaSavedObjectMeta": {
+                    "searchSourceJSON": "{\"index\":\"90943e30-9a47-11e8-b64d-95841ca0b247\",\"filter\":[],\"query\":{\"query\":\"\",\"language\":\"kuery\"}}"
+                }
+            },
+            "references": []
+        },
+        {
+            "id": "some-id",
+            "type": "visualization",
+            "references": [],
+            "attributes": {
+                "title": "Some non-OpenSearch query Vega visualization",
+                "description": "This is a sample Vega visualization without an OpenSearch query.",
+                "kibanaSavedObjectMeta": {},
+                "visState": "{\"title\":\"Vega Visualization without OpenSearch Query\",\"type\":\"vega\",\"params\":{\"spec\":\"{\\\"$schema\\\":\\\"https://vega.github.io/schema/vega-lite/v5.json\\\",\\\"description\\\":\\\"A simple scatter plot.\\\",\\\"data\\\":{\\\"values\\\":[{\\\"x\\\":1,\\\"y\\\":2},{\\\"x\\\":2,\\\"y\\\":3},{\\\"x\\\":3,\\\"y\\\":4},{\\\"x\\\":4,\\\"y\\\":5},{\\\"x\\\":5,\\\"y\\\":6},{\\\"x\\\":6,\\\"y\\\":7},{\\\"x\\\":7,\\\"y\\\":8},{\\\"x\\\":8,\\\"y\\\":9},{\\\"x\\\":9,\\\"y\\\":10}],\\\"mark\\\":\\\"point\\\",\\\"encoding\\\":{\\\"x\\\":{\\\"field\\\":\\\"x\\\",\\\"type\\\":\\\"quantitative\\\"},\\\"y\\\":{\\\"field\\\":\\\"y\\\",\\\"type\\\":\\\"quantitative\\\"}}},\\\"width\\\":400,\\\"height\\\":200}\"},\"aggs\":[]}",
+                "uiStateJSON": "{}",
+                "version": 1
+            }
+        },
+        {
+            "attributes": {
+                "description": "",
+                "kibanaSavedObjectMeta": {
+                    "searchSourceJSON": "{\"query\":{\"query\":\"\",\"language\":\"kuery\"},\"filter\":[]}"
+                },
+                "title": "(Vega) Top destination count",
+                "uiStateJSON": "{}",
+                "version": 1,
+                "visState": "{\"title\":\"(Vega) Top destination count\",\"type\":\"vega\",\"aggs\":[],\"params\":{\"spec\":\"{\\n  $schema: https://vega.github.io/schema/vega-lite/v5.json\\n  data: {\\n    url: {\\n      %context%: true\\n      %timefield%: @timestamp\\n      index: opensearch_dashboards_sample_data_logs\\n      body: {\\n        aggs: {\\n          1: {\\n            terms: {\\n              field: geo.dest\\n              order: {\\n                _count: desc\\n              }\\n              size: 5\\n            }\\n          }\\n        }\\n      }\\n    }\\n    format: {\\n      property: aggregations.1.buckets\\n    }\\n  }\\n  transform: [\\n    {\\n      calculate: datum.key\\n      as: dest\\n    }\\n    {\\n      calculate: datum.doc_count\\n      as: count\\n    }\\n  ]\\n  layer: [\\n    {\\n      mark: {\\n        type: bar\\n        tooltip: true\\n      }\\n    }\\n  ]\\n  encoding: {\\n    x: {\\n      field: count\\n      type: quantitative\\n      axis: {\\n        title: Count\\n      }\\n    }\\n    y: {\\n      field: dest\\n      type: nominal\\n      axis: {\\n        title: Dest\\n      }\\n      sort: -x\\n    }\\n  }\\n}\"}}"
+            },
+            "id": "f0d162c0-227b-11ee-b88b-47a93b5c527c",
+            "migrationVersion": {
+                "visualization": "7.10.0"
+            },
+            "references": [],
+            "type": "visualization",
+            "updated_at": "2023-07-25T19:39:50.773Z",
+            "version": "WzgxLDFd"
+        },
+        {
+            "id": "06cf9c40-9ee8-11e7-8711-e7a007dcef99",
+            "type": "visualization",
+            "updated_at": "2018-08-29T13:22:17.617Z",
+            "version": "1",
+            "migrationVersion": {},
+            "attributes": {
+                "title": "[Logs] Visitors Map",
+                "visState": "{\"title\":\"[Logs] Visitors Map\",\"type\":\"vega\",\"aggs\":[],\"params\":{\"spec\":\"{\\n  $schema: https://vega.github.io/schema/vega/v5.json\\n  config: {\\n    kibana: {type: \\\"map\\\", latitude: 30, longitude: -120, zoom: 3}\\n  }\\n  data: [\\n    {\\n      name: table\\n      url: {\\n        index: opensearch_dashboards_sample_data_logs\\n        %context%: true\\n        %timefield%: timestamp\\n        body: {\\n          size: 0\\n          aggs: {\\n            gridSplit: {\\n              geotile_grid: {field: \\\"geo.coordinates\\\", precision: 5, size: 10000}\\n              aggs: {\\n                gridCentroid: {\\n                  geo_centroid: {\\n                    field: \\\"geo.coordinates\\\"\\n                  }\\n                }\\n              }\\n            }\\n          }\\n        }\\n      }\\n      format: {property: \\\"aggregations.gridSplit.buckets\\\"}\\n      transform: [\\n        {\\n          type: geopoint\\n          projection: projection\\n          fields: [\\n            gridCentroid.location.lon\\n            gridCentroid.location.lat\\n          ]\\n        }\\n      ]\\n    }\\n  ]\\n  scales: [\\n    {\\n      name: gridSize\\n      type: linear\\n      domain: {data: \\\"table\\\", field: \\\"doc_count\\\"}\\n      range: [\\n        50\\n        1000\\n      ]\\n    }\\n    {\\n      name: bubbleColor\\n      type: linear\\n      domain: {\\n        data: table\\n        field: doc_count\\n      }\\n      range: [\\\"rgb(255, 255, 255)\\\",\\\"rgb(249, 212, 204)\\\",\\\"rgb(238, 170, 156)\\\", \\\"rgb(223, 129, 110)\\\"]\\n    }\\n  ]\\n  marks: [\\n    {\\n      name: gridMarker\\n      type: symbol\\n      from: {data: \\\"table\\\"}\\n      encode: {\\n        update: {\\n          fill: {\\n            scale: bubbleColor\\n            field: doc_count\\n          }\\n          size: {scale: \\\"gridSize\\\", field: \\\"doc_count\\\"}\\n          xc: {signal: \\\"datum.x\\\"}\\n          yc: {signal: \\\"datum.y\\\"}\\n          tooltip: {\\n            signal: \\\"{flights: datum.doc_count}\\\"\\n          }\\n        }\\n      }\\n    }\\n  ]\\n}\"}}",
+                "uiStateJSON": "{}",
+                "description": "",
+                "version": 1,
+                "kibanaSavedObjectMeta": {
+                    "searchSourceJSON": "{\"index\":\"90943e30-9a47-11e8-b64d-95841ca0b247\",\"filter\":[],\"query\":{\"query\":\"\",\"language\":\"kuery\"}}"
+                }
+            },
+            "references": []
+        }
+    ]
+}
diff --git a/src/plugins/home/server/services/sample_data/data_sets/util.test.ts b/src/plugins/home/server/services/sample_data/data_sets/util.test.ts
@@ -0,0 +1,65 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { getSavedObjectsWithDataSource } from './util';
+import { SavedObject, updateDataSourceNameInVegaSpec } from '../../../../../../core/server';
+import visualizationObjects from './test_utils/visualization_objects.json';
+
+describe('getSavedObjectsWithDataSource()', () => {
+  const getVisualizationSavedObjects = (): Array<SavedObject<any>> => {
+    // @ts-expect-error
+    return visualizationObjects.saved_objects;
+  };
+
+  test('when processing Vega Visualization saved objects, it should attach data_source_name to each OpenSearch query', () => {
+    const dataSourceId = 'some-datasource-id';
+    const dataSourceName = 'Data Source Name';
+    const expectedUpdatedFields = getVisualizationSavedObjects().map((object) => {
+      const visState = JSON.parse(object.attributes.visState);
+      if (visState.type !== 'vega') {
+        return {
+          vegaSpec: undefined,
+          references: object.references,
+        };
+      }
+      const spec = visState.params.spec;
+      return {
+        vegaSpec: updateDataSourceNameInVegaSpec({
+          newDataSourceName: dataSourceName,
+          spec,
+          spacing: 1,
+        }),
+        references: [
+          {
+            id: dataSourceId,
+            type: 'data-source',
+            name: 'dataSource',
+          },
+        ],
+      };
+    });
+    const updatedVegaVisualizationsFields = getSavedObjectsWithDataSource(
+      getVisualizationSavedObjects(),
+      dataSourceId,
+      dataSourceName
+    ).map((object) => {
+      // @ts-expect-error
+      const visState = JSON.parse(object.attributes.visState);
+      if (visState.type !== 'vega') {
+        return {
+          vegaSpec: undefined,
+          references: object.references,
+        };
+      }
+      const spec = visState.params.spec;
+      return {
+        vegaSpec: spec,
+        references: object.references,
+      };
+    });
+
+    expect(updatedVegaVisualizationsFields).toEqual(expect.arrayContaining(expectedUpdatedFields));
+  });
+});
diff --git a/src/plugins/home/server/services/sample_data/data_sets/util.ts b/src/plugins/home/server/services/sample_data/data_sets/util.ts
@@ -4,6 +4,10 @@
  */
 
 import { SavedObject } from 'opensearch-dashboards/server';
+import {
+  extractVegaSpecFromSavedObject,
+  updateDataSourceNameInVegaSpec,
+} from '../../../../../../core/server';
 
 export const appendDataSourceId = (id: string) => {
   return (dataSourceId?: string) => (dataSourceId ? `${dataSourceId}_` + id : id);
@@ -74,6 +78,31 @@ export const getSavedObjectsWithDataSource = (
         ) {
           saveObject.attributes.title = saveObject.attributes.title + `_${dataSourceTitle}`;
         }
+
+        if (saveObject.type === 'visualization') {
+          const vegaSpec = extractVegaSpecFromSavedObject(saveObject);
+
+          if (!!vegaSpec) {
+            const updatedVegaSpec = updateDataSourceNameInVegaSpec({
+              spec: vegaSpec,
+              newDataSourceName: dataSourceTitle,
+              // Spacing of 1 prevents the Sankey visualization in logs data from exceeding the default url length and breaking
+              spacing: 1,
+            });
+
+            // @ts-expect-error
+            const visStateObject = JSON.parse(saveObject.attributes?.visState);
+            visStateObject.params.spec = updatedVegaSpec;
+
+            // @ts-expect-error
+            saveObject.attributes.visState = JSON.stringify(visStateObject);
+            saveObject.references.push({
+              id: `${dataSourceId}`,
+              type: 'data-source',
+              name: 'dataSource',
+            });
+          }
+        }
       }
 
       return saveObject;