fix: set correct schema on config import (#16041)

* fix: set correct schema on config import

* Fix lint

* Fix test

* Fix tests

* Fix another test

* Fix another test

* Fix base test

* Add helper function

* Fix examples

* Fix test

* Fix test

* Fixing more tests

superset/commands/importers/v1/examples.py

@@ -42,7 +42,7 @@
 from superset.datasets.commands.importers.v1.utils import import_dataset
 from superset.datasets.schemas import ImportV1DatasetSchema
 from superset.models.dashboard import dashboard_slices
-from superset.utils.core import get_example_database
+from superset.utils.core import get_example_database, get_example_default_schema
 
 
 class ImportExamplesCommand(ImportModelsCommand):
@@ -85,7 +85,7 @@ def _get_uuids(cls) -> Set[str]:
         )
 
     @staticmethod
-    def _import(  # pylint: disable=arguments-differ,too-many-locals
+    def _import(  # pylint: disable=arguments-differ, too-many-locals, too-many-branches
         session: Session,
         configs: Dict[str, Any],
         overwrite: bool = False,
@@ -114,6 +114,10 @@ def _import(  # pylint: disable=arguments-differ,too-many-locals
                 else:
                     config["database_id"] = database_ids[config["database_uuid"]]
 
+                # set schema
+                if config["schema"] is None:
+                    config["schema"] = get_example_default_schema()
+
                 dataset = import_dataset(
                     session, config, overwrite=overwrite, force_data=force_data
                 )

superset/connectors/sqla/models.py

@@ -1710,7 +1710,7 @@ def before_update(
         target: "SqlaTable",
     ) -> None:
         """
-        Check whether before update if the target table already exists.
+        Check before update if the target table already exists.
 
         Note this listener is called when any fields are being updated and thus it is
         necessary to first check whether the reference table is being updated.

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

@@ -25,6 +25,7 @@
 from flask import current_app, g
 from sqlalchemy import BigInteger, Boolean, Date, DateTime, Float, String, Text
 from sqlalchemy.orm import Session
+from sqlalchemy.orm.exc import MultipleResultsFound
 from sqlalchemy.sql.visitors import VisitableType
 
 from superset.connectors.sqla.models import SqlaTable
@@ -110,7 +111,19 @@ def import_dataset(
     data_uri = config.get("data")
 
     # import recursively to include columns and metrics
-    dataset = SqlaTable.import_from_dict(session, config, recursive=True, sync=sync)
+    try:
+        dataset = SqlaTable.import_from_dict(session, config, recursive=True, sync=sync)
+    except MultipleResultsFound:
+        # Finding multiple results when importing a dataset only happens because initially
+        # datasets were imported without schemas (eg, `examples.NULL.users`), and later
+        # they were fixed to have the default schema (eg, `examples.public.users`). If a
+        # user created `examples.public.users` during that time the second import will
+        # fail because the UUID match will try to update `examples.NULL.users` to
+        # `examples.public.users`, resulting in a conflict.
+        #
+        # When that happens, we return the original dataset, unmodified.
+        dataset = session.query(SqlaTable).filter_by(uuid=config["uuid"]).one()
+
     if dataset.id is None:
         session.flush()
 

superset/examples/bart_lines.py

@@ -18,7 +18,7 @@
 
 import pandas as pd
 import polyline
-from sqlalchemy import String, Text
+from sqlalchemy import inspect, String, Text
 
 from superset import db
 from superset.utils.core import get_example_database
@@ -29,6 +29,8 @@
 def load_bart_lines(only_metadata: bool = False, force: bool = False) -> None:
     tbl_name = "bart_lines"
     database = get_example_database()
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
     table_exists = database.has_table_by_name(tbl_name)
 
     if not only_metadata and (not table_exists or force):
@@ -40,7 +42,8 @@ def load_bart_lines(only_metadata: bool = False, force: bool = False) -> None:
 
         df.to_sql(
             tbl_name,
-            database.get_sqla_engine(),
+            engine,
+            schema=schema,
             if_exists="replace",
             chunksize=500,
             dtype={
@@ -56,7 +59,7 @@ def load_bart_lines(only_metadata: bool = False, force: bool = False) -> None:
     table = get_table_connector_registry()
     tbl = db.session.query(table).filter_by(table_name=tbl_name).first()
     if not tbl:
-        tbl = table(table_name=tbl_name)
+        tbl = table(table_name=tbl_name, schema=schema)
     tbl.description = "BART lines"
     tbl.database = database
     tbl.filter_select_enabled = True

superset/examples/birth_names.py

@@ -20,12 +20,11 @@
 
 import pandas as pd
 from flask_appbuilder.security.sqla.models import User
-from sqlalchemy import DateTime, String
+from sqlalchemy import DateTime, inspect, String
 from sqlalchemy.sql import column
 
 from superset import app, db, security_manager
-from superset.connectors.base.models import BaseDatasource
-from superset.connectors.sqla.models import SqlMetric, TableColumn
+from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn
 from superset.exceptions import NoDataException
 from superset.models.core import Database
 from superset.models.dashboard import Dashboard
@@ -75,9 +74,13 @@ def load_data(tbl_name: str, database: Database, sample: bool = False) -> None:
         pdf.ds = pd.to_datetime(pdf.ds, unit="ms")
     pdf = pdf.head(100) if sample else pdf
 
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
+
     pdf.to_sql(
         tbl_name,
         database.get_sqla_engine(),
+        schema=schema,
         if_exists="replace",
         chunksize=500,
         dtype={
@@ -98,18 +101,21 @@ def load_birth_names(
     only_metadata: bool = False, force: bool = False, sample: bool = False
 ) -> None:
     """Loading birth name dataset from a zip file in the repo"""
-    tbl_name = "birth_names"
     database = get_example_database()
-    table_exists = database.has_table_by_name(tbl_name)
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
+
+    tbl_name = "birth_names"
+    table_exists = database.has_table_by_name(tbl_name, schema=schema)
 
     if not only_metadata and (not table_exists or force):
         load_data(tbl_name, database, sample=sample)
 
     table = get_table_connector_registry()
-    obj = db.session.query(table).filter_by(table_name=tbl_name).first()
+    obj = db.session.query(table).filter_by(table_name=tbl_name, schema=schema).first()
     if not obj:
         print(f"Creating table [{tbl_name}] reference")
-        obj = table(table_name=tbl_name)
+        obj = table(table_name=tbl_name, schema=schema)
         db.session.add(obj)
 
     _set_table_metadata(obj, database)
@@ -121,14 +127,14 @@ def load_birth_names(
     create_dashboard(slices)
 
 
-def _set_table_metadata(datasource: "BaseDatasource", database: "Database") -> None:
-    datasource.main_dttm_col = "ds"  # type: ignore
+def _set_table_metadata(datasource: SqlaTable, database: "Database") -> None:
+    datasource.main_dttm_col = "ds"
     datasource.database = database
     datasource.filter_select_enabled = True
     datasource.fetch_metadata()
 
 
-def _add_table_metrics(datasource: "BaseDatasource") -> None:
+def _add_table_metrics(datasource: SqlaTable) -> None:
     if not any(col.column_name == "num_california" for col in datasource.columns):
         col_state = str(column("state").compile(db.engine))
         col_num = str(column("num").compile(db.engine))
@@ -147,13 +153,11 @@ def _add_table_metrics(datasource: "BaseDatasource") -> None:
 
     for col in datasource.columns:
         if col.column_name == "ds":
-            col.is_dttm = True  # type: ignore
+            col.is_dttm = True
             break
 
 
-def create_slices(
-    tbl: BaseDatasource, admin_owner: bool
-) -> Tuple[List[Slice], List[Slice]]:
+def create_slices(tbl: SqlaTable, admin_owner: bool) -> Tuple[List[Slice], List[Slice]]:
     metrics = [
         {
             "expressionType": "SIMPLE",

superset/examples/country_map.py

@@ -17,7 +17,7 @@
 import datetime
 
 import pandas as pd
-from sqlalchemy import BigInteger, Date, String
+from sqlalchemy import BigInteger, Date, inspect, String
 from sqlalchemy.sql import column
 
 from superset import db
@@ -38,6 +38,8 @@ def load_country_map_data(only_metadata: bool = False, force: bool = False) -> N
     """Loading data for map with country map"""
     tbl_name = "birth_france_by_region"
     database = utils.get_example_database()
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
     table_exists = database.has_table_by_name(tbl_name)
 
     if not only_metadata and (not table_exists or force):
@@ -48,7 +50,8 @@ def load_country_map_data(only_metadata: bool = False, force: bool = False) -> N
         data["dttm"] = datetime.datetime.now().date()
         data.to_sql(
             tbl_name,
-            database.get_sqla_engine(),
+            engine,
+            schema=schema,
             if_exists="replace",
             chunksize=500,
             dtype={
@@ -76,7 +79,7 @@ def load_country_map_data(only_metadata: bool = False, force: bool = False) -> N
     table = get_table_connector_registry()
     obj = db.session.query(table).filter_by(table_name=tbl_name).first()
     if not obj:
-        obj = table(table_name=tbl_name)
+        obj = table(table_name=tbl_name, schema=schema)
     obj.main_dttm_col = "dttm"
     obj.database = database
     obj.filter_select_enabled = True

superset/examples/energy.py

@@ -18,7 +18,7 @@
 import textwrap
 
 import pandas as pd
-from sqlalchemy import Float, String
+from sqlalchemy import Float, inspect, String
 from sqlalchemy.sql import column
 
 from superset import db
@@ -40,6 +40,8 @@ def load_energy(
     """Loads an energy related dataset to use with sankey and graphs"""
     tbl_name = "energy_usage"
     database = utils.get_example_database()
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
     table_exists = database.has_table_by_name(tbl_name)
 
     if not only_metadata and (not table_exists or force):
@@ -48,7 +50,8 @@ def load_energy(
         pdf = pdf.head(100) if sample else pdf
         pdf.to_sql(
             tbl_name,
-            database.get_sqla_engine(),
+            engine,
+            schema=schema,
             if_exists="replace",
             chunksize=500,
             dtype={"source": String(255), "target": String(255), "value": Float()},
@@ -60,7 +63,7 @@ def load_energy(
     table = get_table_connector_registry()
     tbl = db.session.query(table).filter_by(table_name=tbl_name).first()
     if not tbl:
-        tbl = table(table_name=tbl_name)
+        tbl = table(table_name=tbl_name, schema=schema)
     tbl.description = "Energy consumption"
     tbl.database = database
     tbl.filter_select_enabled = True

superset/examples/flights.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import pandas as pd
-from sqlalchemy import DateTime
+from sqlalchemy import DateTime, inspect
 
 from superset import db
 from superset.utils import core as utils
@@ -27,6 +27,8 @@ def load_flights(only_metadata: bool = False, force: bool = False) -> None:
     """Loading random time series data from a zip file in the repo"""
     tbl_name = "flights"
     database = utils.get_example_database()
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
     table_exists = database.has_table_by_name(tbl_name)
 
     if not only_metadata and (not table_exists or force):
@@ -47,7 +49,8 @@ def load_flights(only_metadata: bool = False, force: bool = False) -> None:
         pdf = pdf.join(airports, on="DESTINATION_AIRPORT", rsuffix="_DEST")
         pdf.to_sql(
             tbl_name,
-            database.get_sqla_engine(),
+            engine,
+            schema=schema,
             if_exists="replace",
             chunksize=500,
             dtype={"ds": DateTime},
@@ -57,7 +60,7 @@ def load_flights(only_metadata: bool = False, force: bool = False) -> None:
     table = get_table_connector_registry()
     tbl = db.session.query(table).filter_by(table_name=tbl_name).first()
     if not tbl:
-        tbl = table(table_name=tbl_name)
+        tbl = table(table_name=tbl_name, schema=schema)
     tbl.description = "Random set of flights in the US"
     tbl.database = database
     tbl.filter_select_enabled = True

superset/examples/long_lat.py

@@ -19,7 +19,7 @@
 
 import geohash
 import pandas as pd
-from sqlalchemy import DateTime, Float, String
+from sqlalchemy import DateTime, Float, inspect, String
 
 from superset import db
 from superset.models.slice import Slice
@@ -38,6 +38,8 @@ def load_long_lat_data(only_metadata: bool = False, force: bool = False) -> None
     """Loading lat/long data from a csv file in the repo"""
     tbl_name = "long_lat"
     database = utils.get_example_database()
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
     table_exists = database.has_table_by_name(tbl_name)
 
     if not only_metadata and (not table_exists or force):
@@ -56,7 +58,8 @@ def load_long_lat_data(only_metadata: bool = False, force: bool = False) -> None
         pdf["delimited"] = pdf["LAT"].map(str).str.cat(pdf["LON"].map(str), sep=",")
         pdf.to_sql(
             tbl_name,
-            database.get_sqla_engine(),
+            engine,
+            schema=schema,
             if_exists="replace",
             chunksize=500,
             dtype={
@@ -85,7 +88,7 @@ def load_long_lat_data(only_metadata: bool = False, force: bool = False) -> None
     table = get_table_connector_registry()
     obj = db.session.query(table).filter_by(table_name=tbl_name).first()
     if not obj:
-        obj = table(table_name=tbl_name)
+        obj = table(table_name=tbl_name, schema=schema)
     obj.main_dttm_col = "datetime"
     obj.database = database
     obj.filter_select_enabled = True

superset/examples/multiformat_time_series.py

@@ -17,7 +17,7 @@
 from typing import Dict, Optional, Tuple
 
 import pandas as pd
-from sqlalchemy import BigInteger, Date, DateTime, String
+from sqlalchemy import BigInteger, Date, DateTime, inspect, String
 
 from superset import app, db
 from superset.models.slice import Slice
@@ -38,6 +38,8 @@ def load_multiformat_time_series(  # pylint: disable=too-many-locals
     """Loading time series data from a zip file in the repo"""
     tbl_name = "multiformat_time_series"
     database = get_example_database()
+    engine = database.get_sqla_engine()
+    schema = inspect(engine).default_schema_name
     table_exists = database.has_table_by_name(tbl_name)
 
     if not only_metadata and (not table_exists or force):
@@ -55,7 +57,8 @@ def load_multiformat_time_series(  # pylint: disable=too-many-locals
 
         pdf.to_sql(
             tbl_name,
-            database.get_sqla_engine(),
+            engine,
+            schema=schema,
             if_exists="replace",
             chunksize=500,
             dtype={
@@ -77,7 +80,7 @@ def load_multiformat_time_series(  # pylint: disable=too-many-locals
     table = get_table_connector_registry()
     obj = db.session.query(table).filter_by(table_name=tbl_name).first()
     if not obj:
-        obj = table(table_name=tbl_name)
+        obj = table(table_name=tbl_name, schema=schema)
     obj.main_dttm_col = "ds"
     obj.database = database
     obj.filter_select_enabled = True