From a69e9deb958256c661434d02404aaec3e18f0af3 Mon Sep 17 00:00:00 2001
From: "Michael S. Molina" <michael.s.molina@gmail.com>
Date: Mon, 24 Jun 2024 15:59:50 -0300
Subject: [PATCH] Adds the migration processor

---
 superset/cli/viz_migrations.py                   |  3 +++
 .../configs/charts/Featured Charts/Sankey.yaml   | 16 ++++++++++++++++
 .../migrations/shared/migrate_viz/processors.py  | 12 ++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/superset/cli/viz_migrations.py b/superset/cli/viz_migrations.py
index 9714969073..6b6ba1a704 100644
--- a/superset/cli/viz_migrations.py
+++ b/superset/cli/viz_migrations.py
@@ -33,6 +33,7 @@ class VizType(str, Enum):
     HISTOGRAM = "histogram"
     LINE = "line"
     PIVOT_TABLE = "pivot_table"
+    SANKEY = "sankey"
     SUNBURST = "sunburst"
     TREEMAP = "treemap"
 
@@ -89,6 +90,7 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
         MigrateHistogramChart,
         MigrateLineChart,
         MigratePivotTable,
+        MigrateSankey,
         MigrateSunburst,
         MigrateTreeMap,
     )
@@ -103,6 +105,7 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
         VizType.HISTOGRAM: MigrateHistogramChart,
         VizType.LINE: MigrateLineChart,
         VizType.PIVOT_TABLE: MigratePivotTable,
+        VizType.SANKEY: MigrateSankey,
         VizType.SUNBURST: MigrateSunburst,
         VizType.TREEMAP: MigrateTreeMap,
     }
diff --git a/superset/examples/configs/charts/Featured Charts/Sankey.yaml b/superset/examples/configs/charts/Featured Charts/Sankey.yaml
index 65fcd2d888..8540f5f649 100644
--- a/superset/examples/configs/charts/Featured Charts/Sankey.yaml	
+++ b/superset/examples/configs/charts/Featured Charts/Sankey.yaml	
@@ -1,3 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
 slice_name: Sankey
 description: null
 certified_by: null
diff --git a/superset/migrations/shared/migrate_viz/processors.py b/superset/migrations/shared/migrate_viz/processors.py
index ef5a1c6bd6..99c545d98f 100644
--- a/superset/migrations/shared/migrate_viz/processors.py
+++ b/superset/migrations/shared/migrate_viz/processors.py
@@ -303,3 +303,15 @@ class MigrateHistogramChart(MigrateViz):
         groupby = self.data.get("groupby")
         if not groupby:
             self.data["groupby"] = []
+
+
+class MigrateSankey(MigrateViz):
+    source_viz_type = "sankey"
+    target_viz_type = "sankey_v2"
+    remove_keys = {"groupby"}
+
+    def _pre_action(self) -> None:
+        groupby = self.data.get("groupby")
+        if groupby and len(groupby) > 1:
+            self.data["source"] = groupby[0]
+            self.data["target"] = groupby[1]