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

docker/.env

@@ -59,7 +59,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

superset/charts/schemas.py

@@ -1550,6 +1550,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):

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 @@
         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 = 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
 
+    _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 @@
 
         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])

superset/commands/chart/importers/v1/__init__.py

@@ -15,21 +15,24 @@
 # 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
 
+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.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):
@@ -46,7 +49,14 @@
     import_error = ChartImportError
 
     @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 datasets associated with charts
         dataset_uuids: set[str] = set()
         for file_name, config in configs.items():
@@ -98,4 +108,12 @@
                 if "query_context" in config:
                     config["query_context"] = None
 
-                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:
+                        new_tag_names = config["tags"]
+                        import_tag(
+                            new_tag_names, contents, chart.id, "chart", db.session
+                        )

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__)
 
@@ -159,6 +162,11 @@
 
         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
 
@@ -173,7 +181,14 @@
 
         if export_related:
             chart_ids = [chart.id for chart in model.slices]
+            dashboard_ids = model.id
+            ExportChartsCommand.disable_tag_export()
             yield from ExportChartsCommand(chart_ids).run()
+            ExportChartsCommand.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,

superset/commands/dashboard/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
@@ -34,10 +34,12 @@
 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.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
 
@@ -57,9 +59,15 @@
     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:
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: Optional[dict[str, Any]] = None,
+    ) -> None:
+        if contents is None:
+            contents = {}
         # discover charts and datasets associated with dashboards
         chart_uuids: set[str] = set()
         dataset_uuids: set[str] = set()
@@ -123,6 +131,14 @@
                 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:
+                        new_tag_names = config["tags"]
+                        import_tag(
+                            new_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])
@@ -143,6 +159,18 @@
                     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:
+                        new_tag_names = config["tags"]
+                        import_tag(
+                            new_tag_names,
+                            contents,
+                            dashboard.id,
+                            "dashboard",
+                            db.session,
+                        )
+
         # set ref in the dashboard_slices table
         values = [
             {"dashboard_id": dashboard_id, "slice_id": chart_id}

superset/commands/database/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,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: Optional[dict[str, Any]] = None,
+    ) -> None:
         # first import databases
         database_ids: dict[str, int] = {}
         for file_name, config in configs.items():

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 @@
     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():

superset/commands/export/assets.py

@@ -27,6 +27,7 @@
 from superset.commands.database.export import ExportDatabasesCommand
 from superset.commands.dataset.export import ExportDatasetsCommand
 from superset.commands.query.export import ExportSavedQueriesCommand
+from superset.commands.tag.export import ExportTagsCommand
 from superset.utils.dict_import_export import EXPORT_VERSION
 
 METADATA_FILE_NAME = "metadata.yaml"
@@ -50,15 +51,17 @@ def run(self) -> Iterator[tuple[str, Callable[[], str]]]:
             ExportDatabasesCommand,
             ExportDatasetsCommand,
             ExportChartsCommand,
+            ExportTagsCommand,
             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():
-                if file_name not in seen:
-                    yield file_name, file_content
-                    seen.add(file_name)
+            if hasattr(command, "dao"):
+                ids = [model.id for model in command.dao.find_all()]
+                for file_name, file_content in command(ids, export_related=False).run():
+                    if file_name not in seen:
+                        yield file_name, file_content
+                        seen.add(file_name)
 
     def validate(self) -> None:
         pass

superset/commands/importers/v1/__init__.py

@@ -61,7 +61,13 @@ 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:
+    def _import(
+        configs: dict[str, Any],
+        overwrite: bool = False,
+        contents: Optional[
+            dict[str, Any]
+        ] = None,  # Use dict[str, str] to match the parent class signature
+    ) -> None:
         raise NotImplementedError("Subclasses MUST implement _import")
 
     @classmethod
@@ -73,7 +79,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:

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 _import(  # pylint: disable=too-many-locals, too-many-branches
         configs: dict[str, Any],
         overwrite: bool = False,
+        contents: Optional[dict[str, Any]] = None,
         force_data: bool = False,
     ) -> None:
         # import databases
@@ -129,7 +130,7 @@
                     dataset = import_dataset(
                         config,
                         overwrite=overwrite,
-                        force_data=force_data,
+                        force_data=bool(force_data),
                         ignore_permissions=True,
                     )
                 except MultipleResultsFound: