feat: API for asset sync (#19220)

* feat: API for asset sync

* Add unit tests.

* Improve tests

* Move files

* Add more tests
This commit is contained in:
Beto Dealmeida 2022-03-22 08:29:24 -07:00 committed by GitHub
parent d3ce398448
commit b05e7dbf2a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 490 additions and 14 deletions

View File

@ -26,3 +26,8 @@ class IncorrectVersionError(CommandException):
class NoValidFilesFoundError(CommandException):
status = 400
message = "No valid import files were found"
class IncorrectFormatError(CommandException):
status = 422
message = "File has the incorrect format"

View File

@ -0,0 +1,163 @@
# 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.
import json
from datetime import datetime
from io import BytesIO
from zipfile import is_zipfile, ZipFile
from flask import request, Response, send_file
from flask_appbuilder.api import BaseApi, expose, protect
from superset.commands.export.assets import ExportAssetsCommand
from superset.commands.importers.exceptions import (
IncorrectFormatError,
NoValidFilesFoundError,
)
from superset.commands.importers.v1.assets import ImportAssetsCommand
from superset.commands.importers.v1.utils import get_contents_from_bundle
from superset.extensions import event_logger
from superset.views.base_api import requires_form_data
class ImportExportRestApi(BaseApi):
"""
API for exporting all assets or importing them.
"""
resource_name = "assets"
openapi_spec_tag = "Import/export"
@expose("/export/", methods=["GET"])
@protect()
@event_logger.log_this_with_context(
action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.export",
log_to_statsd=False,
)
def export(self) -> Response:
"""
Export all assets.
---
get:
description: >-
Returns a ZIP file with all the Superset assets (databases, datasets, charts,
dashboards, saved queries) as YAML files.
responses:
200:
description: ZIP file
content:
application/zip:
schema:
type: string
format: binary
400:
$ref: '#/components/responses/400'
401:
$ref: '#/components/responses/401'
404:
$ref: '#/components/responses/404'
500:
$ref: '#/components/responses/500'
"""
timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
root = f"assets_export_{timestamp}"
filename = f"{root}.zip"
buf = BytesIO()
with ZipFile(buf, "w") as bundle:
for file_name, file_content in ExportAssetsCommand().run():
with bundle.open(f"{root}/{file_name}", "w") as fp:
fp.write(file_content.encode())
buf.seek(0)
response = send_file(
buf,
mimetype="application/zip",
as_attachment=True,
attachment_filename=filename,
)
return response
@expose("/import/", methods=["POST"])
@protect()
@event_logger.log_this_with_context(
action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.import_",
log_to_statsd=False,
)
@requires_form_data
def import_(self) -> Response:
"""Import multiple assets
---
post:
requestBody:
required: true
content:
multipart/form-data:
schema:
type: object
properties:
bundle:
description: upload file (ZIP or JSON)
type: string
format: binary
passwords:
description: >-
JSON map of passwords for each featured database in the
ZIP file. If the ZIP includes a database config in the path
`databases/MyDatabase.yaml`, the password should be provided
in the following format:
`{"databases/MyDatabase.yaml": "my_password"}`.
type: string
responses:
200:
description: Dashboard import result
content:
application/json:
schema:
type: object
properties:
message:
type: string
400:
$ref: '#/components/responses/400'
401:
$ref: '#/components/responses/401'
422:
$ref: '#/components/responses/422'
500:
$ref: '#/components/responses/500'
"""
upload = request.files.get("bundle")
if not upload:
return self.response_400()
if not is_zipfile(upload):
raise IncorrectFormatError("Not a ZIP file")
with ZipFile(upload) as bundle:
contents = get_contents_from_bundle(bundle)
if not contents:
raise NoValidFilesFoundError()
passwords = (
json.loads(request.form["passwords"])
if "passwords" in request.form
else None
)
command = ImportAssetsCommand(contents, passwords=passwords)
command.run()
return self.response(200, message="OK")

View File

@ -67,7 +67,7 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods
self.config = app.config
self.manifest: Dict[Any, Any] = {}
@deprecated(details="use self.superset_app instead of self.flask_app") # type: ignore # pylint: disable=line-too-long,useless-suppression
@deprecated(details="use self.superset_app instead of self.flask_app") # type: ignore
@property
def flask_app(self) -> SupersetApp:
return self.superset_app
@ -143,6 +143,7 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods
from superset.datasets.metrics.api import DatasetMetricRestApi
from superset.explore.form_data.api import ExploreFormDataRestApi
from superset.explore.permalink.api import ExplorePermalinkRestApi
from superset.importexport.api import ImportExportRestApi
from superset.queries.api import QueryRestApi
from superset.queries.saved_queries.api import SavedQueryRestApi
from superset.reports.api import ReportScheduleRestApi
@ -219,6 +220,7 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods
appbuilder.add_api(ExploreFormDataRestApi)
appbuilder.add_api(ExplorePermalinkRestApi)
appbuilder.add_api(FilterSetRestApi)
appbuilder.add_api(ImportExportRestApi)
appbuilder.add_api(QueryRestApi)
appbuilder.add_api(ReportScheduleRestApi)
appbuilder.add_api(ReportExecutionLogRestApi)

View File

@ -14,9 +14,10 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=redefined-outer-name
# pylint: disable=redefined-outer-name, import-outside-toplevel
from typing import Iterator
import importlib
from typing import Any, Iterator
import pytest
from pytest_mock import MockFixture
@ -25,11 +26,12 @@ from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.session import Session
from superset.app import SupersetApp
from superset.extensions import appbuilder
from superset.initialization import SupersetAppInitializer
@pytest.fixture()
def session() -> Iterator[Session]:
@pytest.fixture
def session(mocker: MockFixture) -> Iterator[Session]:
"""
Create an in-memory SQLite session to test models.
"""
@ -40,11 +42,18 @@ def session() -> Iterator[Session]:
# flask calls session.remove()
in_memory_session.remove = lambda: None
# patch session
mocker.patch(
"superset.security.SupersetSecurityManager.get_session",
return_value=in_memory_session,
)
mocker.patch("superset.db.session", in_memory_session)
yield in_memory_session
@pytest.fixture
def app(mocker: MockFixture, session: Session) -> Iterator[SupersetApp]:
@pytest.fixture(scope="module")
def app() -> Iterator[SupersetApp]:
"""
A fixture that generates a Superset app.
"""
@ -52,20 +61,31 @@ def app(mocker: MockFixture, session: Session) -> Iterator[SupersetApp]:
app.config.from_object("superset.config")
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite://"
app.config["FAB_ADD_SECURITY_VIEWS"] = False
app.config["TESTING"] = True
app_initializer = app.config.get("APP_INITIALIZER", SupersetAppInitializer)(app)
# ``superset.extensions.appbuilder`` is a singleton, and won't rebuild the
# routes when this fixture is called multiple times; we need to clear the
# registered views to ensure the initialization can happen more than once.
appbuilder.baseviews = []
app_initializer = SupersetAppInitializer(app)
app_initializer.init_app()
# patch session
mocker.patch(
"superset.security.SupersetSecurityManager.get_session", return_value=session,
)
mocker.patch("superset.db.session", session)
# reload base views to ensure error handlers are applied to the app
with app.app_context():
import superset.views.base
importlib.reload(superset.views.base)
yield app
@pytest.fixture
def client(app: SupersetApp) -> Any:
with app.test_client() as client:
yield client
@pytest.fixture
def app_context(app: SupersetApp) -> Iterator[None]:
"""

View File

@ -0,0 +1,16 @@
# 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.

View File

@ -0,0 +1,254 @@
# 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.
# pylint: disable=invalid-name, import-outside-toplevel
import json
from io import BytesIO
from pathlib import Path
from typing import Any
from zipfile import is_zipfile, ZipFile
from pytest_mock import MockFixture
from superset import security_manager
def test_export_assets(mocker: MockFixture, client: Any) -> None:
"""
Test exporting assets.
"""
from superset.commands.importers.v1.utils import get_contents_from_bundle
# grant access
mocker.patch(
"flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
)
mocker.patch.object(security_manager, "has_access", return_value=True)
mocked_contents = [
(
"metadata.yaml",
"version: 1.0.0\ntype: assets\ntimestamp: '2022-01-01T00:00:00+00:00'\n",
),
("databases/example.yaml", "<DATABASE CONTENTS>"),
]
ExportAssetsCommand = mocker.patch("superset.importexport.api.ExportAssetsCommand")
ExportAssetsCommand().run.return_value = mocked_contents[:]
response = client.get("/api/v1/assets/export/")
assert response.status_code == 200
buf = BytesIO(response.data)
assert is_zipfile(buf)
buf.seek(0)
with ZipFile(buf) as bundle:
contents = get_contents_from_bundle(bundle)
assert contents == dict(mocked_contents)
def test_import_assets(mocker: MockFixture, client: Any) -> None:
"""
Test importing assets.
"""
# grant access
mocker.patch(
"flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
)
mocker.patch.object(security_manager, "has_access", return_value=True)
mocked_contents = {
"metadata.yaml": (
"version: 1.0.0\ntype: assets\ntimestamp: '2022-01-01T00:00:00+00:00'\n"
),
"databases/example.yaml": "<DATABASE CONTENTS>",
}
ImportAssetsCommand = mocker.patch("superset.importexport.api.ImportAssetsCommand")
root = Path("assets_export")
buf = BytesIO()
with ZipFile(buf, "w") as bundle:
for path, contents in mocked_contents.items():
with bundle.open(str(root / path), "w") as fp:
fp.write(contents.encode())
buf.seek(0)
form_data = {
"bundle": (buf, "assets_export.zip"),
"passwords": json.dumps(
{"assets_export/databases/imported_database.yaml": "SECRET"}
),
}
response = client.post(
"/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
)
assert response.status_code == 200
assert response.json == {"message": "OK"}
passwords = {"assets_export/databases/imported_database.yaml": "SECRET"}
ImportAssetsCommand.assert_called_with(mocked_contents, passwords=passwords)
def test_import_assets_not_zip(mocker: MockFixture, client: Any) -> None:
"""
Test error message when the upload is not a ZIP file.
"""
# grant access
mocker.patch(
"flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
)
mocker.patch.object(security_manager, "has_access", return_value=True)
buf = BytesIO(b"definitely_not_a_zip_file")
form_data = {
"bundle": (buf, "broken.txt"),
}
response = client.post(
"/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
)
assert response.status_code == 422
assert response.json == {
"errors": [
{
"message": "Not a ZIP file",
"error_type": "GENERIC_COMMAND_ERROR",
"level": "warning",
"extra": {
"issue_codes": [
{
"code": 1010,
"message": (
"Issue 1010 - Superset encountered an error while "
"running a command."
),
}
]
},
}
]
}
def test_import_assets_no_form_data(mocker: MockFixture, client: Any) -> None:
"""
Test error message when the upload has no form data.
"""
# grant access
mocker.patch(
"flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
)
mocker.patch.object(security_manager, "has_access", return_value=True)
response = client.post("/api/v1/assets/import/", data="some_content")
assert response.status_code == 400
assert response.json == {
"errors": [
{
"message": "Request MIME type is not 'multipart/form-data'",
"error_type": "INVALID_PAYLOAD_FORMAT_ERROR",
"level": "error",
"extra": {
"issue_codes": [
{
"code": 1019,
"message": (
"Issue 1019 - The submitted payload has the incorrect "
"format."
),
}
]
},
}
]
}
def test_import_assets_incorrect_form_data(mocker: MockFixture, client: Any) -> None:
"""
Test error message when the upload form data has the wrong key.
"""
# grant access
mocker.patch(
"flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
)
mocker.patch.object(security_manager, "has_access", return_value=True)
buf = BytesIO(b"definitely_not_a_zip_file")
form_data = {
"wrong": (buf, "broken.txt"),
}
response = client.post(
"/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
)
assert response.status_code == 400
assert response.json == {"message": "Arguments are not correct"}
def test_import_assets_no_contents(mocker: MockFixture, client: Any) -> None:
"""
Test error message when the ZIP bundle has no contents.
"""
# grant access
mocker.patch(
"flask_appbuilder.security.decorators.verify_jwt_in_request", return_value=True
)
mocker.patch.object(security_manager, "has_access", return_value=True)
mocked_contents = {
"README.txt": "Something is wrong",
}
root = Path("assets_export")
buf = BytesIO()
with ZipFile(buf, "w") as bundle:
for path, contents in mocked_contents.items():
with bundle.open(str(root / path), "w") as fp:
fp.write(contents.encode())
buf.seek(0)
form_data = {
"bundle": (buf, "assets_export.zip"),
"passwords": json.dumps(
{"assets_export/databases/imported_database.yaml": "SECRET"}
),
}
response = client.post(
"/api/v1/assets/import/", data=form_data, content_type="multipart/form-data"
)
assert response.status_code == 400
assert response.json == {
"errors": [
{
"message": "No valid import files were found",
"error_type": "GENERIC_COMMAND_ERROR",
"level": "warning",
"extra": {
"issue_codes": [
{
"code": 1010,
"message": (
"Issue 1010 - Superset encountered an error while "
"running a command."
),
}
]
},
}
]
}

View File

@ -0,0 +1,16 @@
# 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.