Skip to content

Commit

Permalink
feat: API for asset sync (#19220)
Browse files Browse the repository at this point in the history
* feat: API for asset sync

* Add unit tests.

* Improve tests

* Move files

* Add more tests
  • Loading branch information
betodealmeida committed Mar 22, 2022
1 parent d3ce398 commit b05e7db
Show file tree
Hide file tree
Showing 7 changed files with 490 additions and 14 deletions.
5 changes: 5 additions & 0 deletions superset/commands/importers/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
163 changes: 163 additions & 0 deletions superset/importexport/api.py
Original file line number Diff line number Diff line change
@@ -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")
4 changes: 3 additions & 1 deletion superset/initialization/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def __init__(self, app: SupersetApp) -> None:
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
Expand Down Expand Up @@ -143,6 +143,7 @@ def init_views(self) -> None:
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
Expand Down Expand Up @@ -219,6 +220,7 @@ def init_views(self) -> None:
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)
Expand Down
46 changes: 33 additions & 13 deletions tests/unit_tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -25,11 +26,12 @@
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.
"""
Expand All @@ -40,32 +42,50 @@ 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.
"""
app = SupersetApp(__name__)

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]:
"""
Expand Down
16 changes: 16 additions & 0 deletions tests/unit_tests/importexport/__init__.py
Original file line number Diff line number Diff line change
@@ -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.
Loading

0 comments on commit b05e7db

Please sign in to comment.