fix: column extra in import/export (apache#17738)

superset/datasets/commands/export.py

@@ -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)

superset/datasets/commands/importers/v1/utils.py

@@ -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__)
 
@@ -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 []
@@ -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(
@@ -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)
@@ -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"):
@@ -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,

superset/datasets/schemas.py

@@ -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)

tests/unit_tests/conftest.py

@@ -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

tests/unit_tests/datasets/__init__.py

@@ -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.

tests/unit_tests/datasets/commands/__init__.py

@@ -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.