From 2ade4f537599c52bea2402460529c694dcbf6227 Mon Sep 17 00:00:00 2001
From: Uladzislau Lasitsa <Uladzislau_Lasitsa@epam.com>
Date: Tue, 9 Mar 2021 10:56:11 +0300
Subject: [PATCH] Migration script for 'schema' (#93427)

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>
---
 .../visualization_migrations.test.ts          | 141 ++++++++++++++++++
 .../saved_objects/visualization_migrations.ts |  46 +++++-
 2 files changed, 185 insertions(+), 2 deletions(-)

diff --git a/src/plugins/visualizations/server/saved_objects/visualization_migrations.test.ts b/src/plugins/visualizations/server/saved_objects/visualization_migrations.test.ts
index 0ca2339f5ab19..5c74de3e2ef9a 100644
--- a/src/plugins/visualizations/server/saved_objects/visualization_migrations.test.ts
+++ b/src/plugins/visualizations/server/saved_objects/visualization_migrations.test.ts
@@ -1762,4 +1762,145 @@ describe('migration visualization', () => {
       });
     });
   });
+
+  describe('7.12.0 update "schema" in aggregations', () => {
+    const migrate = (doc: any) =>
+      visualizationSavedObjectTypeMigrations['7.12.0'](
+        doc as Parameters<SavedObjectMigrationFn>[0],
+        savedObjectMigrationContext
+      );
+    const testDoc = {
+      attributes: {
+        title: 'My Vis',
+        description: 'This is my super cool vis.',
+        visState: JSON.stringify({
+          type: 'metric',
+          title: '[Flights] Delay Type',
+          aggs: [
+            {
+              id: '1',
+              type: 'avg_bucket',
+              schema: 'metric',
+              customBucket: {
+                id: '1-bucket',
+                params: {
+                  orderAgg: {
+                    schema: {
+                      name: 'orderAgg',
+                    },
+                  },
+                },
+                schema: {
+                  name: 'bucketAgg',
+                },
+              },
+              customMetric: {
+                id: '1-metric',
+                schema: {
+                  name: 'metricAgg',
+                },
+                params: {
+                  customMetric: {
+                    schema: {
+                      name: 'metricAgg',
+                    },
+                  },
+                },
+              },
+            },
+            {
+              id: '2',
+              type: 'avg_bucket',
+              schema: 'metric',
+              customBucket: {
+                id: '2-bucket',
+                params: {
+                  orderAgg: {
+                    schema: {
+                      name: 'orderAgg',
+                    },
+                  },
+                },
+              },
+              customMetric: {
+                id: '2-metric',
+                schema: 'metricAgg',
+                params: {
+                  customMetric: {
+                    schema: {
+                      name: 'metricAgg',
+                    },
+                  },
+                },
+              },
+            },
+          ],
+        }),
+      },
+    };
+
+    const expectedDoc = {
+      attributes: {
+        title: 'My Vis',
+        description: 'This is my super cool vis.',
+        visState: JSON.stringify({
+          type: 'metric',
+          title: '[Flights] Delay Type',
+          aggs: [
+            {
+              id: '1',
+              type: 'avg_bucket',
+              schema: 'metric',
+              customBucket: {
+                id: '1-bucket',
+                params: {
+                  orderAgg: {
+                    schema: 'orderAgg',
+                  },
+                },
+                schema: 'bucketAgg',
+              },
+              customMetric: {
+                id: '1-metric',
+                schema: 'metricAgg',
+                params: {
+                  customMetric: {
+                    schema: 'metricAgg',
+                  },
+                },
+              },
+            },
+            {
+              id: '2',
+              type: 'avg_bucket',
+              schema: 'metric',
+              customBucket: {
+                id: '2-bucket',
+                params: {
+                  orderAgg: {
+                    schema: 'orderAgg',
+                  },
+                },
+              },
+              customMetric: {
+                id: '2-metric',
+                schema: 'metricAgg',
+                params: {
+                  customMetric: {
+                    schema: 'metricAgg',
+                  },
+                },
+              },
+            },
+          ],
+        }),
+      },
+    };
+
+    it('should replace all schema object with schema name', () => {
+      const migratedTestDoc = migrate(testDoc);
+
+      expect(migratedTestDoc).toEqual(expectedDoc);
+    });
+  });
 });
diff --git a/src/plugins/visualizations/server/saved_objects/visualization_migrations.ts b/src/plugins/visualizations/server/saved_objects/visualization_migrations.ts
index 89df7aec8d8e6..2c72208b8b283 100644
--- a/src/plugins/visualizations/server/saved_objects/visualization_migrations.ts
+++ b/src/plugins/visualizations/server/saved_objects/visualization_migrations.ts
@@ -6,7 +6,7 @@
  * Side Public License, v 1.
  */
 
-import { cloneDeep, get, omit, has, flow } from 'lodash';
+import { cloneDeep, get, omit, has, flow, forOwn } from 'lodash';
 
 import { SavedObjectMigrationFn } from 'kibana/server';
 
@@ -248,6 +248,48 @@ const migrateDateHistogramAggregation: SavedObjectMigrationFn<any, any> = (doc)
   return doc;
 };
 
+// Migrate schemas inside aggregation (replace 'schema' object to name of the schema)
+const migrateSchema: SavedObjectMigrationFn<any, any> = (doc) => {
+  const visStateJSON = get(doc, 'attributes.visState');
+  let visState;
+
+  if (visStateJSON) {
+    try {
+      visState = JSON.parse(visStateJSON);
+    } catch (e) {
+      // Let it go, the data is invalid and we'll leave it as is
+    }
+
+    function replaceSchema(agg: any) {
+      forOwn(agg, (value: any, key: string) => {
+        if (typeof value === 'object') {
+          if (key === 'schema') {
+            agg[key] = value.name;
+          } else {
+            replaceSchema(value);
+          }
+        }
+      });
+    }
+
+    if (visState && visState.aggs) {
+      for (const agg of visState.aggs) {
+        if (typeof agg === 'object') {
+          replaceSchema(agg);
+        }
+      }
+      return {
+        ...doc,
+        attributes: {
+          ...doc.attributes,
+          visState: JSON.stringify(visState),
+        },
+      };
+    }
+  }
+  return doc;
+};
+
 const removeDateHistogramTimeZones: SavedObjectMigrationFn<any, any> = (doc) => {
   const visStateJSON = get(doc, 'attributes.visState');
   if (visStateJSON) {
@@ -883,5 +925,5 @@ export const visualizationSavedObjectTypeMigrations = {
   '7.9.3': flow(migrateMatchAllQuery),
   '7.10.0': flow(migrateFilterRatioQuery, removeTSVBSearchSource),
   '7.11.0': flow(enableDataTableVisToolbar),
-  '7.12.0': flow(migrateVislibAreaLineBarTypes),
+  '7.12.0': flow(migrateVislibAreaLineBarTypes, migrateSchema),
 };