diff --git a/superset/connectors/druid/views.py b/superset/connectors/druid/views.py index 9947c45911..0f2fcd42f6 100644 --- a/superset/connectors/druid/views.py +++ b/superset/connectors/druid/views.py @@ -233,6 +233,8 @@ class DruidClusterModelView(SupersetModelView, DeleteMixin, YamlExportMixin): ), } + yaml_dict_key = "databases" + edit_form_extra_fields = { "cluster_name": QuerySelectField( "Cluster", diff --git a/superset/utils/import_datasource.py b/superset/utils/import_datasource.py index e924f7f92e..e55862bbfe 100644 --- a/superset/utils/import_datasource.py +++ b/superset/utils/import_datasource.py @@ -26,7 +26,7 @@ def import_datasource( """Imports the datasource from the object to the database. Metrics and columns and datasource will be overrided if exists. - This function can be used to import/export dashboards between multiple + This function can be used to import/export datasources between multiple superset instances. Audit metadata isn't copies over. """ make_transient(i_datasource) diff --git a/superset/views/base.py b/superset/views/base.py index 07346d14f1..5c6e25905c 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -19,7 +19,7 @@ import functools import logging import traceback from datetime import datetime -from typing import Any, Dict +from typing import Any, Dict, Optional import simplejson as json import yaml @@ -208,12 +208,20 @@ def validate_json(form, field): class YamlExportMixin(object): + yaml_dict_key: Optional[str] = None + """ + Override this if you want a dict response instead, with a certain key. + Used on DatabaseView for cli compatibility + """ + @action("yaml_export", __("Export to YAML"), __("Export to YAML?"), "fa-download") def yaml_export(self, items): if not isinstance(items, list): items = [items] data = [t.export_to_dict() for t in items] + if self.yaml_dict_key: + data = {self.yaml_dict_key: data} return Response( yaml.safe_dump(data), headers=generate_download_headers("yaml"), diff --git a/superset/views/database/views.py b/superset/views/database/views.py index be522a266a..150c6bc618 100644 --- a/superset/views/database/views.py +++ b/superset/views/database/views.py @@ -54,6 +54,8 @@ class DatabaseView(DatabaseMixin, SupersetModelView, DeleteMixin, YamlExportMixi edit_template = "superset/models/database/edit.html" validators_columns = {"sqlalchemy_uri": [sqlalchemy_uri_form_validator]} + yaml_dict_key = "databases" + def _delete(self, pk): DeleteMixin._delete(self, pk) diff --git a/tests/dict_import_export_tests.py b/tests/dict_import_export_tests.py index 425eac283f..4890c345c8 100644 --- a/tests/dict_import_export_tests.py +++ b/tests/dict_import_export_tests.py @@ -24,6 +24,7 @@ from superset import db from superset.connectors.druid.models import DruidColumn, DruidDatasource, DruidMetric from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn from superset.utils.core import get_example_database +from superset.utils.dict_import_export import export_to_dict from .base_tests import SupersetTestCase @@ -266,6 +267,26 @@ class DictImportExportTests(SupersetTestCase): imported_copy_table.export_to_dict(), imported_table.export_to_dict() ) + def test_export_datasource_ui_cli(self): + cli_export = export_to_dict( + session=db.session, + recursive=True, + back_references=False, + include_defaults=False, + ) + self.get_resp("/login/", data=dict(username="admin", password="general")) + resp = self.get_resp( + "/databaseview/action_post", {"action": "yaml_export", "rowid": 1} + ) + ui_export = yaml.safe_load(resp) + self.assertEqual( + ui_export["databases"][0]["database_name"], + cli_export["databases"][0]["database_name"], + ) + self.assertEqual( + ui_export["databases"][0]["tables"], cli_export["databases"][0]["tables"] + ) + def test_import_druid_no_metadata(self): datasource, dict_datasource = self.create_druid_datasource( "pure_druid", id=ID_PREFIX + 1