feat(tags): Export and Import Functionality for Superset Dashboards and Charts (#30833)

Co-authored-by: Asher Manangan <[email protected]>

Author: asher-lab

Diff for: docker/.env

@@ -62,7 +62,6 @@ MAPBOX_API_KEY=''
 
 # Make sure you set this to a unique secure random value on production
 SUPERSET_SECRET_KEY=TEST_NON_DEV_SECRET
-
 ENABLE_PLAYWRIGHT=false
 PUPPETEER_SKIP_CHROMIUM_DOWNLOAD=true
 BUILD_SUPERSET_FRONTEND_IN_DOCKER=true

Diff for: superset/charts/schemas.py

@@ -1564,6 +1564,7 @@ class ImportV1ChartSchema(Schema):
     dataset_uuid = fields.UUID(required=True)
     is_managed_externally = fields.Boolean(allow_none=True, dump_default=False)
     external_url = fields.String(allow_none=True)
+    tags = fields.List(fields.String(), allow_none=True)
 
 
 class ChartCacheWarmUpRequestSchema(Schema):

Diff for: superset/commands/chart/export.py

@@ -26,10 +26,13 @@
 from superset.daos.chart import ChartDAO
 from superset.commands.dataset.export import ExportDatasetsCommand
 from superset.commands.export.models import ExportModelsCommand
+from superset.commands.tag.export import ExportTagsCommand
 from superset.models.slice import Slice
+from superset.tags.models import TagType
 from superset.utils.dict_import_export import EXPORT_VERSION
 from superset.utils.file import get_filename
 from superset.utils import json
+from superset.extensions import feature_flag_manager
 
 logger = logging.getLogger(__name__)
 
@@ -71,9 +74,23 @@ def _file_content(model: Slice) -> str:
         if model.table:
             payload["dataset_uuid"] = str(model.table.uuid)
 
+        # Fetch tags from the database if TAGGING_SYSTEM is enabled
+        if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+            tags = getattr(model, "tags", [])
+            payload["tags"] = [tag.name for tag in tags if tag.type == TagType.custom]
         file_content = yaml.safe_dump(payload, sort_keys=False)
         return file_content
 
+    _include_tags: bool = True  # Default to True
+
+    @classmethod
+    def disable_tag_export(cls) -> None:
+        cls._include_tags = False
+
+    @classmethod
+    def enable_tag_export(cls) -> None:
+        cls._include_tags = True
+
     @staticmethod
     def _export(
         model: Slice, export_related: bool = True
@@ -85,3 +102,12 @@ def _export(
 
         if model.table and export_related:
             yield from ExportDatasetsCommand([model.table.id]).run()
+
+        # Check if the calling class is ExportDashboardCommands
+        if (
+            export_related
+            and ExportChartsCommand._include_tags
+            and feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM")
+        ):
+            chart_id = model.id
+            yield from ExportTagsCommand().export(chart_ids=[chart_id])

Diff for: superset/commands/chart/importers/v1/__init__.py

@@ -14,23 +14,27 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
 
 from typing import Any
 
 from marshmallow import Schema
 from sqlalchemy.orm import Session  # noqa: F401
 
+from superset import db
 from superset.charts.schemas import ImportV1ChartSchema
 from superset.commands.chart.exceptions import ChartImportError
 from superset.commands.chart.importers.v1.utils import import_chart
 from superset.commands.database.importers.v1.utils import import_database
 from superset.commands.dataset.importers.v1.utils import import_dataset
 from superset.commands.importers.v1 import ImportModelsCommand
+from superset.commands.importers.v1.utils import import_tag
 from superset.commands.utils import update_chart_config_dataset
 from superset.connectors.sqla.models import SqlaTable
 from superset.daos.chart import ChartDAO
 from superset.databases.schemas import ImportV1DatabaseSchema
 from superset.datasets.schemas import ImportV1DatasetSchema
+from superset.extensions import feature_flag_manager
 
 
 class ImportChartsCommand(ImportModelsCommand):
@@ -47,7 +51,13 @@ class ImportChartsCommand(ImportModelsCommand):
     import_error = ChartImportError
 
     @staticmethod
-    def _import(configs: dict[str, Any], overwrite: bool = False) -> None:  # noqa: C901
+    # ruff: noqa: C901
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: dict[str, Any] | None = None,
+    ) -> None:
+        contents = {} if contents is None else contents
         # discover datasets associated with charts
         dataset_uuids: set[str] = set()
         for file_name, config in configs.items():
@@ -93,4 +103,12 @@ def _import(configs: dict[str, Any], overwrite: bool = False) -> None:  # noqa:
                     "datasource_name": dataset.table_name,
                 }
                 config = update_chart_config_dataset(config, dataset_dict)
-                import_chart(config, overwrite=overwrite)
+                chart = import_chart(config, overwrite=overwrite)
+
+                # Handle tags using import_tag function
+                if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+                    if "tags" in config:
+                        target_tag_names = config["tags"]
+                        import_tag(
+                            target_tag_names, contents, chart.id, "chart", db.session
+                        )

Diff for: superset/commands/dashboard/export.py

@@ -25,6 +25,7 @@
 import yaml
 
 from superset.commands.chart.export import ExportChartsCommand
+from superset.commands.tag.export import ExportTagsCommand
 from superset.commands.dashboard.exceptions import DashboardNotFoundError
 from superset.commands.dashboard.importers.v1.utils import find_chart_uuids
 from superset.daos.dashboard import DashboardDAO
@@ -33,9 +34,11 @@
 from superset.daos.dataset import DatasetDAO
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
+from superset.tags.models import TagType
 from superset.utils.dict_import_export import EXPORT_VERSION
 from superset.utils.file import get_filename
 from superset.utils import json
+from superset.extensions import feature_flag_manager  # Import the feature flag manager
 
 logger = logging.getLogger(__name__)
 
@@ -112,6 +115,7 @@ def _file_name(model: Dashboard) -> str:
         return f"dashboards/{file_name}.yaml"
 
     @staticmethod
+    # ruff: noqa: C901
     def _file_content(model: Dashboard) -> str:
         payload = model.export_to_dict(
             recursive=False,
@@ -159,10 +163,16 @@ def _file_content(model: Dashboard) -> str:
 
         payload["version"] = EXPORT_VERSION
 
+        # Check if the TAGGING_SYSTEM feature is enabled
+        if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+            tags = model.tags if hasattr(model, "tags") else []
+            payload["tags"] = [tag.name for tag in tags if tag.type == TagType.custom]
+
         file_content = yaml.safe_dump(payload, sort_keys=False)
         return file_content
 
     @staticmethod
+    # ruff: noqa: C901
     def _export(
         model: Dashboard, export_related: bool = True
     ) -> Iterator[tuple[str, Callable[[], str]]]:
@@ -173,7 +183,15 @@ def _export(
 
         if export_related:
             chart_ids = [chart.id for chart in model.slices]
-            yield from ExportChartsCommand(chart_ids).run()
+            dashboard_ids = model.id
+            command = ExportChartsCommand(chart_ids)
+            command.disable_tag_export()
+            yield from command.run()
+            command.enable_tag_export()
+            if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+                yield from ExportTagsCommand.export(
+                    dashboard_ids=dashboard_ids, chart_ids=chart_ids
+                )
 
         payload = model.export_to_dict(
             recursive=False,

Diff for: superset/commands/dashboard/importers/v1/__init__.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from __future__ import annotations
+
 from typing import Any
 
 from marshmallow import Schema
@@ -34,11 +36,13 @@
 from superset.commands.database.importers.v1.utils import import_database
 from superset.commands.dataset.importers.v1.utils import import_dataset
 from superset.commands.importers.v1 import ImportModelsCommand
+from superset.commands.importers.v1.utils import import_tag
 from superset.commands.utils import update_chart_config_dataset
 from superset.daos.dashboard import DashboardDAO
 from superset.dashboards.schemas import ImportV1DashboardSchema
 from superset.databases.schemas import ImportV1DatabaseSchema
 from superset.datasets.schemas import ImportV1DatasetSchema
+from superset.extensions import feature_flag_manager
 from superset.migrations.shared.native_filters import migrate_dashboard
 from superset.models.dashboard import Dashboard, dashboard_slices
 
@@ -58,9 +62,15 @@ class ImportDashboardsCommand(ImportModelsCommand):
     import_error = DashboardImportError
 
     # TODO (betodealmeida): refactor to use code from other commands
-    # pylint: disable=too-many-branches, too-many-locals
+    # pylint: disable=too-many-branches, too-many-locals, too-many-statements
     @staticmethod
-    def _import(configs: dict[str, Any], overwrite: bool = False) -> None:  # noqa: C901
+    # ruff: noqa: C901
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: dict[str, Any] | None = None,
+    ) -> None:
+        contents = {} if contents is None else contents
         # discover charts and datasets associated with dashboards
         chart_uuids: set[str] = set()
         dataset_uuids: set[str] = set()
@@ -120,6 +130,14 @@ def _import(configs: dict[str, Any], overwrite: bool = False) -> None:  # noqa:
                 charts.append(chart)
                 chart_ids[str(chart.uuid)] = chart.id
 
+                # Handle tags using import_tag function
+                if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+                    if "tags" in config:
+                        target_tag_names = config["tags"]
+                        import_tag(
+                            target_tag_names, contents, chart.id, "chart", db.session
+                        )
+
         # store the existing relationship between dashboards and charts
         existing_relationships = db.session.execute(
             select([dashboard_slices.c.dashboard_id, dashboard_slices.c.slice_id])
@@ -140,6 +158,18 @@ def _import(configs: dict[str, Any], overwrite: bool = False) -> None:  # noqa:
                     if (dashboard.id, chart_id) not in existing_relationships:
                         dashboard_chart_ids.append((dashboard.id, chart_id))
 
+                # Handle tags using import_tag function
+                if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):
+                    if "tags" in config:
+                        target_tag_names = config["tags"]
+                        import_tag(
+                            target_tag_names,
+                            contents,
+                            dashboard.id,
+                            "dashboard",
+                            db.session,
+                        )
+
         # set ref in the dashboard_slices table
         values = [
             {"dashboard_id": dashboard_id, "slice_id": chart_id}

Diff for: superset/commands/database/importers/v1/__init__.py

@@ -15,6 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from __future__ import annotations
+
 from typing import Any
 
 from marshmallow import Schema
@@ -42,7 +44,11 @@ class ImportDatabasesCommand(ImportModelsCommand):
     import_error = DatabaseImportError
 
     @staticmethod
-    def _import(configs: dict[str, Any], overwrite: bool = False) -> None:
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: dict[str, Any] | None = None,
+    ) -> None:
         # first import databases
         database_ids: dict[str, int] = {}
         for file_name, config in configs.items():

Diff for: superset/commands/dataset/importers/v1/__init__.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import Any
+from typing import Any, Optional
 
 from marshmallow import Schema
 from sqlalchemy.orm import Session  # noqa: F401
@@ -42,7 +42,13 @@ class ImportDatasetsCommand(ImportModelsCommand):
     import_error = DatasetImportError
 
     @staticmethod
-    def _import(configs: dict[str, Any], overwrite: bool = False) -> None:
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: Optional[dict[str, Any]] = None,
+    ) -> None:
+        if contents is None:
+            contents = {}
         # discover databases associated with datasets
         database_uuids: set[str] = set()
         for file_name, config in configs.items():

Diff for: superset/commands/export/assets.py

@@ -53,6 +53,7 @@ def run(self) -> Iterator[tuple[str, Callable[[], str]]]:
             ExportDashboardsCommand,
             ExportSavedQueriesCommand,
         ]
+
         for command in commands:
             ids = [model.id for model in command.dao.find_all()]
             for file_name, file_content in command(ids, export_related=False).run():

Diff for: superset/commands/importers/v1/__init__.py

@@ -14,8 +14,11 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
+from __future__ import annotations
+
 import logging
-from typing import Any, Optional
+from typing import Any
 
 from marshmallow import Schema, validate  # noqa: F401
 from marshmallow.exceptions import ValidationError
@@ -64,7 +67,12 @@ def __init__(self, contents: dict[str, str], *args: Any, **kwargs: Any):
         self._configs: dict[str, Any] = {}
 
     @staticmethod
-    def _import(configs: dict[str, Any], overwrite: bool = False) -> None:
+    # ruff: noqa: C901
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: dict[str, Any] | None = None,
+    ) -> None:
         raise NotImplementedError("Subclasses MUST implement _import")
 
     @classmethod
@@ -76,7 +84,7 @@ def run(self) -> None:
         self.validate()
 
         try:
-            self._import(self._configs, self.overwrite)
+            self._import(self._configs, self.overwrite, self.contents)
         except CommandException:
             raise
         except Exception as ex:
@@ -87,7 +95,7 @@ def validate(self) -> None:  # noqa: F811
 
         # verify that the metadata file is present and valid
         try:
-            metadata: Optional[dict[str, str]] = load_metadata(self.contents)
+            metadata: dict[str, str] | None = load_metadata(self.contents)
         except ValidationError as exc:
             exceptions.append(exc)
             metadata = None

Diff for: superset/commands/importers/v1/examples.py

@@ -14,7 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Any
+from typing import Any, Optional
 
 from marshmallow import Schema
 from sqlalchemy.exc import MultipleResultsFound
@@ -90,6 +90,7 @@ def _get_uuids(cls) -> set[str]:
     def _import(  # pylint: disable=too-many-locals, too-many-branches  # noqa: C901
         configs: dict[str, Any],
         overwrite: bool = False,
+        contents: Optional[dict[str, Any]] = None,
         force_data: bool = False,
     ) -> None:
         # import databases