From 4b49a044c27752e357c5b654b184cc809b98d4d2 Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Tue, 24 Mar 2020 15:21:39 +0000 Subject: [PATCH 1/2] [dataset] New, export API endpoint --- superset/datasets/api.py | 69 +++++++++++++++++++++++++++++++--- superset/datasets/schemas.py | 2 + tests/dataset_api_tests.py | 73 ++++++++++++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 5 deletions(-) diff --git a/superset/datasets/api.py b/superset/datasets/api.py index 2e12629c42f1..f0e1ab8984f7 100644 --- a/superset/datasets/api.py +++ b/superset/datasets/api.py @@ -16,8 +16,9 @@ # under the License. import logging +import yaml from flask import g, request, Response -from flask_appbuilder.api import expose, protect, safe +from flask_appbuilder.api import expose, protect, rison, safe from flask_appbuilder.models.sqla.interface import SQLAInterface from superset.connectors.sqla.models import SqlaTable @@ -33,8 +34,12 @@ DatasetUpdateFailedError, ) from superset.datasets.commands.update import UpdateDatasetCommand -from superset.datasets.schemas import DatasetPostSchema, DatasetPutSchema -from superset.views.base import DatasourceFilter +from superset.datasets.schemas import ( + DatasetPostSchema, + DatasetPutSchema, + get_export_ids_schema, +) +from superset.views.base import DatasourceFilter, generate_download_headers from superset.views.base_api import BaseSupersetModelRestApi from superset.views.database.filters import DatabaseFilter @@ -49,8 +54,10 @@ class DatasetRestApi(BaseSupersetModelRestApi): allow_browser_login = True class_permission_name = "TableModelView" - include_route_methods = RouteMethod.REST_MODEL_VIEW_CRUD_SET | {RouteMethod.RELATED} - + include_route_methods = RouteMethod.REST_MODEL_VIEW_CRUD_SET | { + RouteMethod.EXPORT, + RouteMethod.RELATED, + } list_columns = [ "changed_by_name", "changed_by_url", @@ -268,3 +275,55 @@ def delete(self, pk: int) -> Response: # pylint: disable=arguments-differ except DatasetDeleteFailedError as e: logger.error(f"Error deleting model {self.__class__.__name__}: {e}") return self.response_422(message=str(e)) + + @expose("/export/", methods=["GET"]) + @protect() + @safe + @rison(get_export_ids_schema) + def export(self, **kwargs): + """Export dashboards + --- + get: + description: >- + Exports multiple datasets and downloads them as YAML files + parameters: + - in: query + name: q + content: + application/json: + schema: + type: array + items: + type: integer + responses: + 200: + description: Dataset export + content: + text/plain: + schema: + type: string + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 500: + $ref: '#/components/responses/500' + """ + requested_ids = kwargs["rison"] + query = self.datamodel.session.query(SqlaTable).filter( + SqlaTable.id.in_(requested_ids) + ) + query = self._base_filters.apply_all(query) + items = query.all() + ids = [item.id for item in items] + if len(ids) != len(requested_ids): + return self.response_404() + + data = [t.export_to_dict() for t in items] + return Response( + yaml.safe_dump(data), + headers=generate_download_headers("yaml"), + mimetype="application/text", + ) diff --git a/superset/datasets/schemas.py b/superset/datasets/schemas.py index 370550da619c..28736b3ea3e0 100644 --- a/superset/datasets/schemas.py +++ b/superset/datasets/schemas.py @@ -18,6 +18,8 @@ from marshmallow import fields, Schema from marshmallow.validate import Length +get_export_ids_schema = {"type": "array", "items": {"type": "integer"}} + class DatasetPostSchema(Schema): database = fields.Integer(required=True) diff --git a/tests/dataset_api_tests.py b/tests/dataset_api_tests.py index a55140a5dae7..4392e0abaa6f 100644 --- a/tests/dataset_api_tests.py +++ b/tests/dataset_api_tests.py @@ -20,6 +20,8 @@ from unittest.mock import patch import prison +import yaml +from sqlalchemy.sql import func from superset import db, security_manager from superset.connectors.sqla.models import SqlaTable @@ -30,6 +32,8 @@ ) from superset.models.core import Database from superset.utils.core import get_example_database +from superset.utils.dict_import_export import export_to_dict +from superset.views.base import generate_download_headers from .base_tests import SupersetTestCase @@ -50,6 +54,15 @@ def insert_dataset( db.session.commit() return table + @staticmethod + def get_birth_names_dataset(): + example_db = get_example_database() + return ( + db.session.query(SqlaTable) + .filter_by(database=example_db, table_name="birth_names") + .one() + ) + def test_get_dataset_list(self): """ Dataset API: Test get dataset list @@ -452,3 +465,63 @@ def test_delete_dataset_sqlalchemy_error(self, mock_dao_delete): self.assertEqual(data, {"message": "Dataset could not be deleted."}) db.session.delete(table) db.session.commit() + + def test_export_dataset(self): + """ + Dataset API: Test export dataset + :return: + """ + birth_names_dataset = self.get_birth_names_dataset() + + argument = [birth_names_dataset.id] + uri = f"api/v1/dataset/export/?q={prison.dumps(argument)}" + + self.login(username="admin") + rv = self.client.get(uri) + self.assertEqual(rv.status_code, 200) + self.assertEqual( + rv.headers["Content-Disposition"], + generate_download_headers("yaml")["Content-Disposition"], + ) + + cli_export = export_to_dict( + session=db.session, + recursive=True, + back_references=False, + include_defaults=False, + ) + cli_export_tables = cli_export["databases"][0]["tables"] + expected_response = [] + for export_table in cli_export_tables: + if export_table["table_name"] == "birth_names": + expected_response = export_table + break + ui_export = yaml.safe_load(rv.data.decode("utf-8")) + self.assertEqual(ui_export[0], expected_response) + + def test_export_dataset_not_found(self): + """ + Dataset API: Test export dataset not found + :return: + """ + max_id = db.session.query(func.max(SqlaTable.id)).scalar() + # Just one does not exist and we get 404 + argument = [max_id + 1, 1] + uri = f"api/v1/dataset/export/?q={prison.dumps(argument)}" + self.login(username="admin") + rv = self.client.get(uri) + self.assertEqual(rv.status_code, 404) + + def test_export_dataset_gamma(self): + """ + Dataset API: Test export dataset has gamma + :return: + """ + birth_names_dataset = self.get_birth_names_dataset() + + argument = [birth_names_dataset.id] + uri = f"api/v1/dataset/export/?q={prison.dumps(argument)}" + + self.login(username="gamma") + rv = self.client.get(uri) + self.assertEqual(rv.status_code, 401) From 777152d2ce03cdb8bc8d382aa99c737a2c122c38 Mon Sep 17 00:00:00 2001 From: Daniel Gaspar Date: Tue, 24 Mar 2020 23:05:10 +0000 Subject: [PATCH 2/2] Fix, lint --- tests/dataset_api_tests.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/dataset_api_tests.py b/tests/dataset_api_tests.py index dc3bd672df1c..77e056680c8a 100644 --- a/tests/dataset_api_tests.py +++ b/tests/dataset_api_tests.py @@ -34,7 +34,6 @@ from superset.utils.core import get_example_database from superset.utils.dict_import_export import export_to_dict from superset.views.base import generate_download_headers - from tests.base_tests import SupersetTestCase