Skip to content

Commit

Permalink
fix: column extra in import/export (apache#17738)
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida authored and shcoderAlex committed Feb 7, 2022
1 parent 9e64354 commit c19e601
Show file tree
Hide file tree
Showing 10 changed files with 512 additions and 25 deletions.
15 changes: 9 additions & 6 deletions superset/datasets/commands/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,12 +59,15 @@ def _export(model: SqlaTable) -> Iterator[Tuple[str, str]]:
payload[key] = json.loads(payload[key])
except json.decoder.JSONDecodeError:
logger.info("Unable to decode `%s` field: %s", key, payload[key])
for metric in payload.get("metrics", []):
if metric.get("extra"):
try:
metric["extra"] = json.loads(metric["extra"])
except json.decoder.JSONDecodeError:
logger.info("Unable to decode `extra` field: %s", metric["extra"])
for key in ("metrics", "columns"):
for attributes in payload.get(key, []):
if attributes.get("extra"):
try:
attributes["extra"] = json.loads(attributes["extra"])
except json.decoder.JSONDecodeError:
logger.info(
"Unable to decode `extra` field: %s", attributes["extra"]
)

payload["version"] = EXPORT_VERSION
payload["database_uuid"] = str(model.database.uuid)
Expand Down
32 changes: 16 additions & 16 deletions superset/datasets/commands/importers/v1/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@

from superset.connectors.sqla.models import SqlaTable
from superset.models.core import Database
from superset.utils.core import get_example_database

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -96,13 +95,17 @@ def import_dataset(
config[key] = json.dumps(config[key])
except TypeError:
logger.info("Unable to encode `%s` field: %s", key, config[key])
for metric in config.get("metrics", []):
if metric.get("extra") is not None:
try:
metric["extra"] = json.dumps(metric["extra"])
except TypeError:
logger.info("Unable to encode `extra` field: %s", metric["extra"])
metric["extra"] = None
for key in ("metrics", "columns"):
for attributes in config.get(key, []):
# should be a dictionary, but in initial exports this was a string
if isinstance(attributes.get("extra"), dict):
try:
attributes["extra"] = json.dumps(attributes["extra"])
except TypeError:
logger.info(
"Unable to encode `extra` field: %s", attributes["extra"]
)
attributes["extra"] = None

# should we delete columns and metrics not present in the current import?
sync = ["columns", "metrics"] if overwrite else []
Expand All @@ -127,9 +130,8 @@ def import_dataset(
if dataset.id is None:
session.flush()

example_database = get_example_database()
try:
table_exists = example_database.has_table_by_name(dataset.table_name)
table_exists = dataset.database.has_table_by_name(dataset.table_name)
except Exception: # pylint: disable=broad-except
# MySQL doesn't play nice with GSheets table names
logger.warning(
Expand All @@ -139,7 +141,7 @@ def import_dataset(

if data_uri and (not table_exists or force_data):
logger.info("Downloading data from %s", data_uri)
load_data(data_uri, dataset, example_database, session)
load_data(data_uri, dataset, dataset.database, session)

if hasattr(g, "user") and g.user:
dataset.owners.append(g.user)
Expand All @@ -148,7 +150,7 @@ def import_dataset(


def load_data(
data_uri: str, dataset: SqlaTable, example_database: Database, session: Session
data_uri: str, dataset: SqlaTable, database: Database, session: Session
) -> None:
data = request.urlopen(data_uri) # pylint: disable=consider-using-with
if data_uri.endswith(".gz"):
Expand All @@ -162,14 +164,12 @@ def load_data(
df[column_name] = pd.to_datetime(df[column_name])

# reuse session when loading data if possible, to make import atomic
if example_database.sqlalchemy_uri == current_app.config.get(
"SQLALCHEMY_DATABASE_URI"
) or not current_app.config.get("SQLALCHEMY_EXAMPLES_URI"):
if database.sqlalchemy_uri == current_app.config.get("SQLALCHEMY_DATABASE_URI"):
logger.info("Loading data inside the import transaction")
connection = session.connection()
else:
logger.warning("Loading data outside the import transaction")
connection = example_database.get_sqla_engine()
connection = database.get_sqla_engine()

df.to_sql(
dataset.table_name,
Expand Down
3 changes: 2 additions & 1 deletion superset/datasets/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,8 @@ class DatasetRelatedObjectsResponse(Schema):

class ImportV1ColumnSchema(Schema):
column_name = fields.String(required=True)
extra = fields.Dict(allow_none=True)
# extra was initially exported incorrectly as a string
extra = fields.Raw(allow_none=True)
verbose_name = fields.String(allow_none=True)
is_dttm = fields.Boolean(default=False, allow_none=True)
is_active = fields.Boolean(default=True, allow_none=True)
Expand Down
41 changes: 39 additions & 2 deletions tests/unit_tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,25 +14,62 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=redefined-outer-name

from typing import Iterator

import pytest
from pytest_mock import MockFixture
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm.session import Session

from superset.app import SupersetApp
from superset.initialization import SupersetAppInitializer


@pytest.fixture()
def session() -> Iterator[Session]:
"""
Create an in-memory SQLite session to test models.
"""
engine = create_engine("sqlite://")
Session_ = sessionmaker(bind=engine) # pylint: disable=invalid-name
in_memory_session = Session_()

# flask calls session.remove()
in_memory_session.remove = lambda: None

yield in_memory_session


@pytest.fixture
def app_context():
def app(mocker: MockFixture, session: Session) -> Iterator[SupersetApp]:
"""
A fixture for running the test inside an app context.
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_initializer = app.config.get("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)

yield app


@pytest.fixture
def app_context(app: SupersetApp) -> Iterator[None]:
"""
A fixture that yields and application context.
"""
with app.app_context():
yield
16 changes: 16 additions & 0 deletions tests/unit_tests/datasets/__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.
16 changes: 16 additions & 0 deletions tests/unit_tests/datasets/commands/__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 c19e601

Please sign in to comment.