From e53a50a15594e4eeb4914203d016f43847ac4031 Mon Sep 17 00:00:00 2001 From: ofekisr <35701650+ofekisr@users.noreply.github.com> Date: Sun, 16 Jan 2022 08:32:50 +0200 Subject: [PATCH] chore(supeset.utils.core): move all database utils to database utils module (#18058) --- superset/cli.py | 8 +- superset/commands/importers/v1/examples.py | 3 +- superset/examples/bart_lines.py | 2 +- superset/examples/birth_names.py | 2 +- superset/examples/country_map.py | 4 +- superset/examples/energy.py | 4 +- superset/examples/flights.py | 4 +- superset/examples/long_lat.py | 4 +- superset/examples/multiformat_time_series.py | 2 +- superset/examples/paris.py | 4 +- superset/examples/random_time_series.py | 4 +- superset/examples/sf_population_polygons.py | 4 +- superset/examples/world_bank.py | 3 +- superset/utils/core.py | 47 +----------- superset/utils/database.py | 74 +++++++++++++++++++ superset/utils/mock_data.py | 2 +- tests/integration_tests/access_tests.py | 2 +- tests/integration_tests/alerts_tests.py | 5 +- tests/integration_tests/base_tests.py | 3 +- tests/integration_tests/celery_tests.py | 3 +- .../charts/data/api_tests.py | 3 +- tests/integration_tests/config_tests.py | 2 +- tests/integration_tests/conftest.py | 4 +- tests/integration_tests/core_tests.py | 30 ++++---- .../css_templates/api_tests.py | 2 +- tests/integration_tests/csv_upload_tests.py | 7 +- .../integration_tests/databases/api_tests.py | 2 +- .../databases/commands_tests.py | 3 +- tests/integration_tests/datasets/api_tests.py | 8 +- .../datasets/commands_tests.py | 3 +- tests/integration_tests/datasource_tests.py | 3 +- .../db_engine_specs/base_engine_spec_tests.py | 2 +- .../dict_import_export_tests.py | 2 +- .../fixtures/birth_names_dashboard.py | 3 +- .../integration_tests/fixtures/datasource.py | 3 +- .../fixtures/energy_dashboard.py | 3 +- .../fixtures/unicode_dashboard.py | 3 +- .../fixtures/world_bank_dashboard.py | 3 +- .../integration_tests/import_export_tests.py | 3 +- .../importexport/commands_tests.py | 2 +- .../integration_tests/jinja_context_tests.py | 27 +++---- tests/integration_tests/model_tests.py | 2 +- tests/integration_tests/queries/api_tests.py | 2 +- .../queries/saved_queries/api_tests.py | 2 +- .../queries/saved_queries/commands_tests.py | 2 +- tests/integration_tests/reports/api_tests.py | 2 +- .../reports/commands_tests.py | 2 +- tests/integration_tests/security_tests.py | 2 +- .../integration_tests/sql_validator_tests.py | 2 +- tests/integration_tests/sqla_models_tests.py | 2 +- tests/integration_tests/sqllab_tests.py | 3 +- tests/integration_tests/strategy_tests.py | 2 +- tests/integration_tests/utils_tests.py | 2 +- 53 files changed, 184 insertions(+), 143 deletions(-) create mode 100644 superset/utils/database.py diff --git a/superset/cli.py b/superset/cli.py index 99a9343096..8b55a81409 100755 --- a/superset/cli.py +++ b/superset/cli.py @@ -37,9 +37,9 @@ from flask_appbuilder import Model from flask_appbuilder.api import BaseApi from flask_appbuilder.api.manager import resolver +import superset.utils.database as database_utils from superset import app, appbuilder, config, security_manager from superset.extensions import celery_app, db -from superset.utils import core as utils from superset.utils.celery import session_scope from superset.utils.encrypt import SecretsMigrator from superset.utils.urls import get_url_path @@ -120,7 +120,7 @@ def load_examples_run( if only_metadata: print("Loading examples metadata") else: - examples_db = utils.get_example_database() + examples_db = database_utils.get_example_database() print(f"Loading examples metadata and related data into {examples_db}") # pylint: disable=import-outside-toplevel @@ -227,7 +227,7 @@ def import_directory(directory: str, overwrite: bool, force: bool) -> None: ) def set_database_uri(database_name: str, uri: str, skip_create: bool) -> None: """Updates a database connection URI""" - utils.get_or_create_db(database_name, uri, not skip_create) + database_utils.get_or_create_db(database_name, uri, not skip_create) @superset.command() @@ -768,7 +768,7 @@ def load_test_users_run() -> None: sm = security_manager - examples_db = utils.get_example_database() + examples_db = database_utils.get_example_database() examples_pv = sm.add_permission_view_menu("database_access", examples_db.perm) diff --git a/superset/commands/importers/v1/examples.py b/superset/commands/importers/v1/examples.py index 0fb1ce255d..43d7e7e1fa 100644 --- a/superset/commands/importers/v1/examples.py +++ b/superset/commands/importers/v1/examples.py @@ -42,7 +42,8 @@ from superset.datasets.commands.importers.v1 import ImportDatasetsCommand 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, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database class ImportExamplesCommand(ImportModelsCommand): diff --git a/superset/examples/bart_lines.py b/superset/examples/bart_lines.py index a57275f632..bcdf3589a8 100644 --- a/superset/examples/bart_lines.py +++ b/superset/examples/bart_lines.py @@ -21,8 +21,8 @@ import polyline from sqlalchemy import inspect, String, Text from superset import db -from superset.utils.core import get_example_database +from ..utils.database import get_example_database from .helpers import get_example_data, get_table_connector_registry diff --git a/superset/examples/birth_names.py b/superset/examples/birth_names.py index ef964e298b..722c0e9f44 100644 --- a/superset/examples/birth_names.py +++ b/superset/examples/birth_names.py @@ -29,8 +29,8 @@ from superset.exceptions import NoDataException from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.slice import Slice -from superset.utils.core import get_example_database +from ..utils.database import get_example_database from .helpers import ( get_example_data, get_slice_json, diff --git a/superset/examples/country_map.py b/superset/examples/country_map.py index f35135df2c..049de6650c 100644 --- a/superset/examples/country_map.py +++ b/superset/examples/country_map.py @@ -20,10 +20,10 @@ import pandas as pd from sqlalchemy import BigInteger, Date, inspect, String from sqlalchemy.sql import column +import superset.utils.database as database_utils from superset import db from superset.connectors.sqla.models import SqlMetric from superset.models.slice import Slice -from superset.utils import core as utils from .helpers import ( get_example_data, @@ -37,7 +37,7 @@ from .helpers import ( def load_country_map_data(only_metadata: bool = False, force: bool = False) -> None: """Loading data for map with country map""" tbl_name = "birth_france_by_region" - database = utils.get_example_database() + database = 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) diff --git a/superset/examples/energy.py b/superset/examples/energy.py index 5d74c87ce2..137d7fe735 100644 --- a/superset/examples/energy.py +++ b/superset/examples/energy.py @@ -21,10 +21,10 @@ import pandas as pd from sqlalchemy import Float, inspect, String from sqlalchemy.sql import column +import superset.utils.database as database_utils from superset import db from superset.connectors.sqla.models import SqlMetric from superset.models.slice import Slice -from superset.utils import core as utils from .helpers import ( get_example_data, @@ -39,7 +39,7 @@ def load_energy( ) -> None: """Loads an energy related dataset to use with sankey and graphs""" tbl_name = "energy_usage" - database = utils.get_example_database() + database = 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) diff --git a/superset/examples/flights.py b/superset/examples/flights.py index d38830b463..46fdc5c1d0 100644 --- a/superset/examples/flights.py +++ b/superset/examples/flights.py @@ -17,8 +17,8 @@ import pandas as pd from sqlalchemy import DateTime, inspect +import superset.utils.database as database_utils from superset import db -from superset.utils import core as utils from .helpers import get_example_data, get_table_connector_registry @@ -26,7 +26,7 @@ from .helpers import get_example_data, get_table_connector_registry 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() + database = 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) diff --git a/superset/examples/long_lat.py b/superset/examples/long_lat.py index 1c9b0bcffc..4245be1057 100644 --- a/superset/examples/long_lat.py +++ b/superset/examples/long_lat.py @@ -21,9 +21,9 @@ import geohash import pandas as pd from sqlalchemy import DateTime, Float, inspect, String +import superset.utils.database as database_utils from superset import db from superset.models.slice import Slice -from superset.utils import core as utils from .helpers import ( get_example_data, @@ -37,7 +37,7 @@ from .helpers import ( 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() + database = 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) diff --git a/superset/examples/multiformat_time_series.py b/superset/examples/multiformat_time_series.py index caecbaa904..1e9ee497db 100644 --- a/superset/examples/multiformat_time_series.py +++ b/superset/examples/multiformat_time_series.py @@ -21,8 +21,8 @@ from sqlalchemy import BigInteger, Date, DateTime, inspect, String from superset import app, db from superset.models.slice import Slice -from superset.utils.core import get_example_database +from ..utils.database import get_example_database from .helpers import ( get_example_data, get_slice_json, diff --git a/superset/examples/paris.py b/superset/examples/paris.py index 87d8823513..264d80feeb 100644 --- a/superset/examples/paris.py +++ b/superset/examples/paris.py @@ -19,15 +19,15 @@ import json import pandas as pd from sqlalchemy import inspect, String, Text +import superset.utils.database as database_utils from superset import db -from superset.utils import core as utils from .helpers import get_example_data, get_table_connector_registry def load_paris_iris_geojson(only_metadata: bool = False, force: bool = False) -> None: tbl_name = "paris_iris_mapping" - database = utils.get_example_database() + database = 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) diff --git a/superset/examples/random_time_series.py b/superset/examples/random_time_series.py index 56f9a4f54c..0f39b95bd1 100644 --- a/superset/examples/random_time_series.py +++ b/superset/examples/random_time_series.py @@ -18,9 +18,9 @@ import pandas as pd from sqlalchemy import DateTime, inspect, String +import superset.utils.database as database_utils from superset import app, db from superset.models.slice import Slice -from superset.utils import core as utils from .helpers import ( get_example_data, @@ -35,7 +35,7 @@ def load_random_time_series_data( ) -> None: """Loading random time series data from a zip file in the repo""" tbl_name = "random_time_series" - database = utils.get_example_database() + database = 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) diff --git a/superset/examples/sf_population_polygons.py b/superset/examples/sf_population_polygons.py index c34e61262d..6e60d60121 100644 --- a/superset/examples/sf_population_polygons.py +++ b/superset/examples/sf_population_polygons.py @@ -19,8 +19,8 @@ import json import pandas as pd from sqlalchemy import BigInteger, Float, inspect, Text +import superset.utils.database as database_utils from superset import db -from superset.utils import core as utils from .helpers import get_example_data, get_table_connector_registry @@ -29,7 +29,7 @@ def load_sf_population_polygons( only_metadata: bool = False, force: bool = False ) -> None: tbl_name = "sf_population_polygons" - database = utils.get_example_database() + database = 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) diff --git a/superset/examples/world_bank.py b/superset/examples/world_bank.py index 9d0b6a8aa9..2b3eff9725 100644 --- a/superset/examples/world_bank.py +++ b/superset/examples/world_bank.py @@ -23,6 +23,7 @@ import pandas as pd from sqlalchemy import DateTime, inspect, String from sqlalchemy.sql import column +import superset.utils.database from superset import app, db from superset.connectors.sqla.models import SqlMetric from superset.models.dashboard import Dashboard @@ -46,7 +47,7 @@ def load_world_bank_health_n_pop( # pylint: disable=too-many-locals, too-many-s ) -> None: """Loads the world bank health dataset, slices and a dashboard""" tbl_name = "wb_health_population" - database = utils.get_example_database() + database = superset.utils.database.get_example_database() engine = database.get_sqla_engine() schema = inspect(engine).default_schema_name table_exists = database.has_table_by_name(tbl_name) diff --git a/superset/utils/core.py b/superset/utils/core.py index 002592ae09..4908fd98dc 100644 --- a/superset/utils/core.py +++ b/superset/utils/core.py @@ -86,7 +86,6 @@ from sqlalchemy.types import TEXT, TypeDecorator, TypeEngine from typing_extensions import TypedDict, TypeGuard from superset.constants import ( - EXAMPLES_DB_UUID, EXTRA_FORM_DATA_APPEND_KEYS, EXTRA_FORM_DATA_OVERRIDE_EXTRA_KEYS, EXTRA_FORM_DATA_OVERRIDE_REGULAR_MAPPINGS, @@ -107,6 +106,7 @@ from superset.typing import ( FormData, Metric, ) +from superset.utils.database import get_example_database from superset.utils.dates import datetime_to_epoch, EPOCH from superset.utils.hashing import md5_sha_from_dict, md5_sha_from_str @@ -117,8 +117,6 @@ except ImportError: if TYPE_CHECKING: from superset.connectors.base.models import BaseColumn, BaseDatasource - from superset.models.core import Database - logging.getLogger("MARKDOWN").setLevel(logging.INFO) logger = logging.getLogger(__name__) @@ -1207,49 +1205,6 @@ def user_label(user: User) -> Optional[str]: return None -def get_or_create_db( - database_name: str, sqlalchemy_uri: str, always_create: Optional[bool] = True -) -> "Database": - # pylint: disable=import-outside-toplevel - from superset import db - from superset.models import core as models - - database = ( - db.session.query(models.Database).filter_by(database_name=database_name).first() - ) - - # databases with a fixed UUID - uuids = { - "examples": EXAMPLES_DB_UUID, - } - - if not database and always_create: - logger.info("Creating database reference for %s", database_name) - database = models.Database( - database_name=database_name, uuid=uuids.get(database_name) - ) - db.session.add(database) - - if database: - database.set_sqlalchemy_uri(sqlalchemy_uri) - db.session.commit() - - return database - - -def get_example_database() -> "Database": - db_uri = ( - current_app.config.get("SQLALCHEMY_EXAMPLES_URI") - or current_app.config["SQLALCHEMY_DATABASE_URI"] - ) - return get_or_create_db("examples", db_uri) - - -def get_main_database() -> "Database": - db_uri = current_app.config["SQLALCHEMY_DATABASE_URI"] - return get_or_create_db("main", db_uri) - - def get_example_default_schema() -> Optional[str]: """ Return the default schema of the examples database, if any. diff --git a/superset/utils/database.py b/superset/utils/database.py new file mode 100644 index 0000000000..146412d57c --- /dev/null +++ b/superset/utils/database.py @@ -0,0 +1,74 @@ +# 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. +from __future__ import annotations + +import logging +from typing import Optional, TYPE_CHECKING + +from flask import current_app + +from superset.constants import EXAMPLES_DB_UUID + +if TYPE_CHECKING: + from superset.connectors.sqla.models import Database + +logging.getLogger("MARKDOWN").setLevel(logging.INFO) +logger = logging.getLogger(__name__) + + +# TODO: duplicate code with DatabaseDao, below function should be moved or use dao +def get_or_create_db( + database_name: str, sqlalchemy_uri: str, always_create: Optional[bool] = True +) -> Database: + # pylint: disable=import-outside-toplevel + from superset import db + from superset.models import core as models + + database = ( + db.session.query(models.Database).filter_by(database_name=database_name).first() + ) + + # databases with a fixed UUID + uuids = { + "examples": EXAMPLES_DB_UUID, + } + + if not database and always_create: + logger.info("Creating database reference for %s", database_name) + database = models.Database( + database_name=database_name, uuid=uuids.get(database_name) + ) + db.session.add(database) + + if database: + database.set_sqlalchemy_uri(sqlalchemy_uri) + db.session.commit() + + return database + + +def get_example_database() -> Database: + db_uri = ( + current_app.config.get("SQLALCHEMY_EXAMPLES_URI") + or current_app.config["SQLALCHEMY_DATABASE_URI"] + ) + return get_or_create_db("examples", db_uri) + + +def get_main_database() -> Database: + db_uri = current_app.config["SQLALCHEMY_DATABASE_URI"] + return get_or_create_db("main", db_uri) diff --git a/superset/utils/mock_data.py b/superset/utils/mock_data.py index 52838f0478..e51bb39e20 100644 --- a/superset/utils/mock_data.py +++ b/superset/utils/mock_data.py @@ -181,7 +181,7 @@ def add_data( :param bool append: if the table already exists, append data or replace? """ # pylint: disable=import-outside-toplevel - from superset.utils.core import get_example_database + from superset.utils.database import get_example_database database = get_example_database() table_exists = database.has_table_by_name(table_name) diff --git a/tests/integration_tests/access_tests.py b/tests/integration_tests/access_tests.py index 06a3ed2380..13febbd413 100644 --- a/tests/integration_tests/access_tests.py +++ b/tests/integration_tests/access_tests.py @@ -42,7 +42,7 @@ from superset.connectors.druid.models import DruidDatasource from superset.connectors.sqla.models import SqlaTable from superset.models import core as models from superset.models.datasource_access_request import DatasourceAccessRequest -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from .base_tests import SupersetTestCase diff --git a/tests/integration_tests/alerts_tests.py b/tests/integration_tests/alerts_tests.py index 6a373093d5..50558741d8 100644 --- a/tests/integration_tests/alerts_tests.py +++ b/tests/integration_tests/alerts_tests.py @@ -22,6 +22,7 @@ from unittest.mock import patch import pytest from sqlalchemy.orm import Session +import superset.utils.database from superset import db from superset.exceptions import SupersetException from superset.models.alerts import Alert, AlertLog, SQLObservation @@ -56,7 +57,7 @@ logger = logging.getLogger(__name__) @pytest.yield_fixture(scope="module") def setup_database(): with app.app_context(): - example_database = utils.get_example_database() + example_database = superset.utils.database.get_example_database() example_database.get_sqla_engine().execute( "CREATE TABLE test_table AS SELECT 1 as first, 2 as second" ) @@ -88,7 +89,7 @@ def create_alert( recipients="recipient1@superset.com", slack_channel="#test_channel", sql=sql, - database_id=utils.get_example_database().id, + database_id=superset.utils.database.get_example_database().id, validator_type=validator_type, validator_config=validator_config, ) diff --git a/tests/integration_tests/base_tests.py b/tests/integration_tests/base_tests.py index d77f216c73..fcd5e09088 100644 --- a/tests/integration_tests/base_tests.py +++ b/tests/integration_tests/base_tests.py @@ -45,7 +45,8 @@ from superset.models.slice import Slice from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.datasource_access_request import DatasourceAccessRequest -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from superset.views.base_api import BaseSupersetModelRestApi FAKE_DB_NAME = "fake_db_100" diff --git a/tests/integration_tests/celery_tests.py b/tests/integration_tests/celery_tests.py index 1fd076f25b..f68ad85cb5 100644 --- a/tests/integration_tests/celery_tests.py +++ b/tests/integration_tests/celery_tests.py @@ -44,7 +44,8 @@ from superset.errors import ErrorLevel, SupersetErrorType from superset.extensions import celery_app from superset.models.sql_lab import Query from superset.sql_parse import ParsedQuery, CtasMethod -from superset.utils.core import get_example_database, backend +from superset.utils.core import backend +from superset.utils.database import get_example_database CELERY_SLEEP_TIME = 6 QUERY = "SELECT name FROM birth_names LIMIT 1" diff --git a/tests/integration_tests/charts/data/api_tests.py b/tests/integration_tests/charts/data/api_tests.py index efe751c895..8d9797f67e 100644 --- a/tests/integration_tests/charts/data/api_tests.py +++ b/tests/integration_tests/charts/data/api_tests.py @@ -47,11 +47,10 @@ from superset.models.slice import Slice from superset.typing import AdhocColumn from superset.utils.core import ( AnnotationType, - get_example_database, get_example_default_schema, - get_main_database, AdhocMetricExpressionType, ) +from superset.utils.database import get_example_database, get_main_database from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType from tests.common.query_context_generator import ANNOTATION_LAYERS diff --git a/tests/integration_tests/config_tests.py b/tests/integration_tests/config_tests.py index 6ac76793ca..45528913ef 100644 --- a/tests/integration_tests/config_tests.py +++ b/tests/integration_tests/config_tests.py @@ -24,7 +24,7 @@ from tests.integration_tests.test_app import app from superset import db from superset.connectors.sqla.models import SqlaTable -from superset.utils.core import get_or_create_db +from superset.utils.database import get_or_create_db FULL_DTTM_DEFAULTS_EXAMPLE = { "main_dttm_col": "id", diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index ca10a31a69..ec22df22d4 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -25,8 +25,8 @@ from unittest.mock import patch from tests.integration_tests.test_app import app from superset import db from superset.extensions import feature_flag_manager -from superset.utils.core import get_example_database, json_dumps_w_dates - +from superset.utils.core import json_dumps_w_dates +from superset.utils.database import get_example_database CTAS_SCHEMA_NAME = "sqllab_test_db" ADMIN_SCHEMA_NAME = "admin_database" diff --git a/tests/integration_tests/core_tests.py b/tests/integration_tests/core_tests.py index 67e5e5da0a..1c4682ad9a 100644 --- a/tests/integration_tests/core_tests.py +++ b/tests/integration_tests/core_tests.py @@ -25,6 +25,8 @@ import json import logging from typing import Dict, List from urllib.parse import quote + +import superset.utils.database from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, load_birth_names_data, @@ -41,7 +43,7 @@ import pandas as pd import sqlalchemy as sqla from sqlalchemy.exc import SQLAlchemyError from superset.models.cache import CacheKey -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.conftest import with_feature_flags from tests.integration_tests.fixtures.energy_dashboard import ( load_energy_table_with_slice, @@ -154,7 +156,7 @@ class TestCore(SupersetTestCase): self.assertEqual(cache_key_with_groupby, viz.cache_key(qobj)) def test_get_superset_tables_not_allowed(self): - example_db = utils.get_example_database() + example_db = superset.utils.database.get_example_database() schema_name = self.default_schema_backend_map[example_db.backend] self.login(username="gamma") uri = f"superset/tables/{example_db.id}/{schema_name}/undefined/" @@ -162,7 +164,7 @@ class TestCore(SupersetTestCase): self.assertEqual(rv.status_code, 404) def test_get_superset_tables_substr(self): - example_db = utils.get_example_database() + example_db = superset.utils.database.get_example_database() if example_db.backend in {"presto", "hive"}: # TODO: change table to the real table that is in examples. return @@ -471,7 +473,7 @@ class TestCore(SupersetTestCase): # need to temporarily allow sqlite dbs, teardown will undo this app.config["PREVENT_UNSAFE_DB_CONNECTIONS"] = False self.login(username=username) - database = utils.get_example_database() + database = superset.utils.database.get_example_database() # validate that the endpoint works with the password-masked sqlalchemy uri data = json.dumps( { @@ -560,7 +562,7 @@ class TestCore(SupersetTestCase): self.assertEqual(expected_body, response_body) def test_custom_password_store(self): - database = utils.get_example_database() + database = superset.utils.database.get_example_database() conn_pre = sqla.engine.url.make_url(database.sqlalchemy_uri_decrypted) def custom_password_store(uri): @@ -578,13 +580,13 @@ class TestCore(SupersetTestCase): # validate that sending a password-masked uri does not over-write the decrypted # uri self.login(username=username) - database = utils.get_example_database() + database = superset.utils.database.get_example_database() sqlalchemy_uri_decrypted = database.sqlalchemy_uri_decrypted url = "databaseview/edit/{}".format(database.id) data = {k: database.__getattribute__(k) for k in DatabaseView.add_columns} data["sqlalchemy_uri"] = database.safe_sqlalchemy_uri() self.client.post(url, data=data) - database = utils.get_example_database() + database = superset.utils.database.get_example_database() self.assertEqual(sqlalchemy_uri_decrypted, database.sqlalchemy_uri_decrypted) # Need to clean up after ourselves @@ -737,14 +739,14 @@ class TestCore(SupersetTestCase): @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_extra_table_metadata(self): self.login() - example_db = utils.get_example_database() + example_db = superset.utils.database.get_example_database() schema = "default" if example_db.backend in {"presto", "hive"} else "superset" self.get_json_resp( f"/superset/extra_table_metadata/{example_db.id}/birth_names/{schema}/" ) def test_templated_sql_json(self): - if utils.get_example_database().backend == "presto": + if superset.utils.database.get_example_database().backend == "presto": # TODO: make it work for presto return self.login() @@ -1222,7 +1224,7 @@ class TestCore(SupersetTestCase): @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_select_star(self): self.login(username="admin") - examples_db = utils.get_example_database() + examples_db = superset.utils.database.get_example_database() resp = self.get_resp(f"/superset/select_star/{examples_db.id}/birth_names") self.assertIn("gender", resp) @@ -1231,7 +1233,7 @@ class TestCore(SupersetTestCase): Database API: Test get select star not allowed """ self.login(username="gamma") - example_db = utils.get_example_database() + example_db = superset.utils.database.get_example_database() resp = self.client.get(f"/superset/select_star/{example_db.id}/birth_names") self.assertEqual(resp.status_code, 403) @@ -1467,7 +1469,7 @@ class TestCore(SupersetTestCase): def test_virtual_table_explore_visibility(self): # test that default visibility it set to True - database = utils.get_example_database() + database = superset.utils.database.get_example_database() self.assertEqual(database.allows_virtual_table_explore, True) # test that visibility is disabled when extra is set to False @@ -1489,8 +1491,8 @@ class TestCore(SupersetTestCase): self.assertEqual(database.allows_virtual_table_explore, True) def test_explore_database_id(self): - database = utils.get_example_database() - explore_database = utils.get_example_database() + database = superset.utils.database.get_example_database() + explore_database = superset.utils.database.get_example_database() # test that explore_database_id is the regular database # id if none is set in the extra diff --git a/tests/integration_tests/css_templates/api_tests.py b/tests/integration_tests/css_templates/api_tests.py index 4e714318b3..8f7e580e8c 100644 --- a/tests/integration_tests/css_templates/api_tests.py +++ b/tests/integration_tests/css_templates/api_tests.py @@ -24,7 +24,7 @@ from sqlalchemy.sql import func import tests.integration_tests.test_app from superset import db from superset.models.core import CssTemplate -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase diff --git a/tests/integration_tests/csv_upload_tests.py b/tests/integration_tests/csv_upload_tests.py index 59ffd05e78..499d0c82ee 100644 --- a/tests/integration_tests/csv_upload_tests.py +++ b/tests/integration_tests/csv_upload_tests.py @@ -27,6 +27,7 @@ from unittest import mock import pandas as pd import pytest +import superset.utils.database from superset.sql_parse import Table from tests.integration_tests.conftest import ADMIN_SCHEMA_NAME from tests.integration_tests.test_app import app # isort:skip @@ -61,11 +62,11 @@ def setup_csv_upload(): with app.app_context(): login(test_client, username="admin") - upload_db = utils.get_or_create_db( + upload_db = superset.utils.database.get_or_create_db( CSV_UPLOAD_DATABASE, app.config["SQLALCHEMY_EXAMPLES_URI"] ) extra = upload_db.get_extra() - extra["explore_database_id"] = utils.get_example_database().id + extra["explore_database_id"] = superset.utils.database.get_example_database().id upload_db.extra = json.dumps(extra) upload_db.allow_file_upload = True db.session.commit() @@ -275,7 +276,7 @@ def test_import_csv_explore_database(setup_csv_upload, create_csv_files): resp = upload_csv(CSV_FILENAME1, CSV_UPLOAD_TABLE_W_EXPLORE) assert f'CSV file "{CSV_FILENAME1}" uploaded to table "{full_table_name}"' in resp table = SupersetTestCase.get_table(name=CSV_UPLOAD_TABLE_W_EXPLORE) - assert table.database_id == utils.get_example_database().id + assert table.database_id == superset.utils.database.get_example_database().id @pytest.mark.usefixtures("setup_csv_upload") diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py index a64eb3f7dc..78d7285e01 100644 --- a/tests/integration_tests/databases/api_tests.py +++ b/tests/integration_tests/databases/api_tests.py @@ -43,7 +43,7 @@ from superset.db_engine_specs.hana import HanaEngineSpec from superset.errors import SupersetError from superset.models.core import Database, ConfigurationMethod from superset.models.reports import ReportSchedule, ReportScheduleType -from superset.utils.core import get_example_database, get_main_database +from superset.utils.database import get_example_database, get_main_database from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, diff --git a/tests/integration_tests/databases/commands_tests.py b/tests/integration_tests/databases/commands_tests.py index 3d682f700c..199b08ce10 100644 --- a/tests/integration_tests/databases/commands_tests.py +++ b/tests/integration_tests/databases/commands_tests.py @@ -45,7 +45,8 @@ from superset.exceptions import ( SupersetTimeoutException, ) from superset.models.core import Database -from superset.utils.core import backend, get_example_database +from superset.utils.core import backend +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py index 03209fd94f..0e7606c45e 100644 --- a/tests/integration_tests/datasets/api_tests.py +++ b/tests/integration_tests/datasets/api_tests.py @@ -35,12 +35,8 @@ from superset.dao.exceptions import ( ) from superset.extensions import db, security_manager from superset.models.core import Database -from superset.utils.core import ( - backend, - get_example_database, - get_example_default_schema, - get_main_database, -) +from superset.utils.core import backend, get_example_default_schema +from superset.utils.database import get_example_database, get_main_database from superset.utils.dict_import_export import export_to_dict from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.conftest import CTAS_SCHEMA_NAME diff --git a/tests/integration_tests/datasets/commands_tests.py b/tests/integration_tests/datasets/commands_tests.py index ac0319aa91..af22d9319b 100644 --- a/tests/integration_tests/datasets/commands_tests.py +++ b/tests/integration_tests/datasets/commands_tests.py @@ -30,7 +30,8 @@ from superset.datasets.commands.exceptions import DatasetNotFoundError from superset.datasets.commands.export import ExportDatasetsCommand from superset.datasets.commands.importers import v0, v1 from superset.models.core import Database -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.fixtures.energy_dashboard import ( load_energy_table_data, diff --git a/tests/integration_tests/datasource_tests.py b/tests/integration_tests/datasource_tests.py index c923f2ebbf..eecc0501b3 100644 --- a/tests/integration_tests/datasource_tests.py +++ b/tests/integration_tests/datasource_tests.py @@ -27,7 +27,8 @@ from superset.connectors.sqla.models import SqlaTable from superset.datasets.commands.exceptions import DatasetNotFoundError from superset.exceptions import SupersetGenericDBErrorException from superset.models.core import Database -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import db_insert_temp_object, SupersetTestCase from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, diff --git a/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py b/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py index e568b59a2b..d0e20ccdf5 100644 --- a/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py +++ b/tests/integration_tests/db_engine_specs/base_engine_spec_tests.py @@ -31,7 +31,7 @@ from superset.db_engine_specs.mysql import MySQLEngineSpec from superset.db_engine_specs.sqlite import SqliteEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.sql_parse import ParsedQuery -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec from tests.integration_tests.test_app import app diff --git a/tests/integration_tests/dict_import_export_tests.py b/tests/integration_tests/dict_import_export_tests.py index fe7ff512f1..bd9f79e4ae 100644 --- a/tests/integration_tests/dict_import_export_tests.py +++ b/tests/integration_tests/dict_import_export_tests.py @@ -31,7 +31,7 @@ from superset.connectors.druid.models import ( DruidCluster, ) from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from superset.utils.dict_import_export import export_to_dict from .base_tests import SupersetTestCase diff --git a/tests/integration_tests/fixtures/birth_names_dashboard.py b/tests/integration_tests/fixtures/birth_names_dashboard.py index 89c3c8161a..7d9616f57f 100644 --- a/tests/integration_tests/fixtures/birth_names_dashboard.py +++ b/tests/integration_tests/fixtures/birth_names_dashboard.py @@ -26,7 +26,8 @@ from superset.connectors.sqla.models import SqlaTable from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.slice import Slice -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.common.example_data_generator.birth_names.birth_names_generator_factory import ( BirthNamesGeneratorFactory, ) diff --git a/tests/integration_tests/fixtures/datasource.py b/tests/integration_tests/fixtures/datasource.py index 86ab6cf153..b6f2476f66 100644 --- a/tests/integration_tests/fixtures/datasource.py +++ b/tests/integration_tests/fixtures/datasource.py @@ -17,7 +17,8 @@ """Fixtures for test_datasource.py""" from typing import Any, Dict -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database def get_datasource_post() -> Dict[str, Any]: diff --git a/tests/integration_tests/fixtures/energy_dashboard.py b/tests/integration_tests/fixtures/energy_dashboard.py index f4362ea746..c0291db2a9 100644 --- a/tests/integration_tests/fixtures/energy_dashboard.py +++ b/tests/integration_tests/fixtures/energy_dashboard.py @@ -24,7 +24,8 @@ from sqlalchemy import column, Float, String from superset import db from superset.connectors.sqla.models import SqlaTable, SqlMetric from superset.models.slice import Slice -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.integration_tests.dashboard_utils import create_slice, create_table_metadata from tests.integration_tests.test_app import app diff --git a/tests/integration_tests/fixtures/unicode_dashboard.py b/tests/integration_tests/fixtures/unicode_dashboard.py index 1f1218c8f7..c7b828176f 100644 --- a/tests/integration_tests/fixtures/unicode_dashboard.py +++ b/tests/integration_tests/fixtures/unicode_dashboard.py @@ -22,7 +22,8 @@ from superset import db from superset.connectors.sqla.models import SqlaTable from superset.models.dashboard import Dashboard from superset.models.slice import Slice -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.integration_tests.dashboard_utils import ( create_dashboard, create_slice, diff --git a/tests/integration_tests/fixtures/world_bank_dashboard.py b/tests/integration_tests/fixtures/world_bank_dashboard.py index 0782c8bdf5..1ac1706a9d 100644 --- a/tests/integration_tests/fixtures/world_bank_dashboard.py +++ b/tests/integration_tests/fixtures/world_bank_dashboard.py @@ -29,7 +29,8 @@ from superset.connectors.sqla.models import SqlaTable from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.slice import Slice -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.integration_tests.dashboard_utils import ( create_dashboard, create_table_metadata, diff --git a/tests/integration_tests/import_export_tests.py b/tests/integration_tests/import_export_tests.py index 1745f5e01c..f34cf62150 100644 --- a/tests/integration_tests/import_export_tests.py +++ b/tests/integration_tests/import_export_tests.py @@ -45,7 +45,8 @@ from superset.dashboards.commands.importers.v0 import import_chart, import_dashb from superset.datasets.commands.importers.v0 import import_dataset from superset.models.dashboard import Dashboard from superset.models.slice import Slice -from superset.utils.core import get_example_database, get_example_default_schema +from superset.utils.core import get_example_default_schema +from superset.utils.database import get_example_database from tests.integration_tests.fixtures.world_bank_dashboard import ( load_world_bank_dashboard_with_slices, diff --git a/tests/integration_tests/importexport/commands_tests.py b/tests/integration_tests/importexport/commands_tests.py index 230462d097..ceaf097565 100644 --- a/tests/integration_tests/importexport/commands_tests.py +++ b/tests/integration_tests/importexport/commands_tests.py @@ -22,7 +22,7 @@ from freezegun import freeze_time from superset import security_manager from superset.databases.commands.export import ExportDatabasesCommand -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase diff --git a/tests/integration_tests/jinja_context_tests.py b/tests/integration_tests/jinja_context_tests.py index 31f877740b..924e93e17e 100644 --- a/tests/integration_tests/jinja_context_tests.py +++ b/tests/integration_tests/jinja_context_tests.py @@ -22,6 +22,7 @@ from unittest import mock import pytest from sqlalchemy.dialects.postgresql import dialect +import superset.utils.database import tests.integration_tests.test_app from superset import app from superset.exceptions import SupersetTemplateException @@ -264,77 +265,77 @@ class TestJinja2Context(SupersetTestCase): safe_proxy(func, {"foo": lambda: "bar"}) def test_process_template(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() sql = "SELECT '{{ 1+1 }}'" tp = get_template_processor(database=maindb) rendered = tp.process_template(sql) self.assertEqual("SELECT '2'", rendered) def test_get_template_kwarg(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo }}" tp = get_template_processor(database=maindb, foo="bar") rendered = tp.process_template(s) self.assertEqual("bar", rendered) def test_template_kwarg(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo }}" tp = get_template_processor(database=maindb) rendered = tp.process_template(s, foo="bar") self.assertEqual("bar", rendered) def test_get_template_kwarg_dict(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo.bar }}" tp = get_template_processor(database=maindb, foo={"bar": "baz"}) rendered = tp.process_template(s) self.assertEqual("baz", rendered) def test_template_kwarg_dict(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo.bar }}" tp = get_template_processor(database=maindb) rendered = tp.process_template(s, foo={"bar": "baz"}) self.assertEqual("baz", rendered) def test_get_template_kwarg_lambda(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo() }}" tp = get_template_processor(database=maindb, foo=lambda: "bar") with pytest.raises(SupersetTemplateException): tp.process_template(s) def test_template_kwarg_lambda(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo() }}" tp = get_template_processor(database=maindb) with pytest.raises(SupersetTemplateException): tp.process_template(s, foo=lambda: "bar") def test_get_template_kwarg_module(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ dt(2017, 1, 1).isoformat() }}" tp = get_template_processor(database=maindb, dt=datetime) with pytest.raises(SupersetTemplateException): tp.process_template(s) def test_template_kwarg_module(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ dt(2017, 1, 1).isoformat() }}" tp = get_template_processor(database=maindb) with pytest.raises(SupersetTemplateException): tp.process_template(s, dt=datetime) def test_get_template_kwarg_nested_module(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo.dt }}" tp = get_template_processor(database=maindb, foo={"dt": datetime}) with pytest.raises(SupersetTemplateException): tp.process_template(s) def test_template_kwarg_nested_module(self) -> None: - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "{{ foo.dt }}" tp = get_template_processor(database=maindb) with pytest.raises(SupersetTemplateException): @@ -353,7 +354,7 @@ class TestJinja2Context(SupersetTestCase): @mock.patch("superset.jinja_context.context_addons") def test_template_context_addons(self, addons_mock) -> None: addons_mock.return_value = {"datetime": datetime} - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() s = "SELECT '{{ datetime(2017, 1, 1).isoformat() }}'" tp = get_template_processor(database=maindb) rendered = tp.process_template(s) @@ -414,7 +415,7 @@ class TestJinja2Context(SupersetTestCase): def test_custom_template_processors_ignored(self) -> None: """Test custom template processor is ignored for a difference backend database.""" - maindb = utils.get_example_database() + maindb = superset.utils.database.get_example_database() sql = "SELECT '$DATE()'" tp = get_template_processor(database=maindb) rendered = tp.process_template(sql) diff --git a/tests/integration_tests/model_tests.py b/tests/integration_tests/model_tests.py index bcd57726ff..6371a06123 100644 --- a/tests/integration_tests/model_tests.py +++ b/tests/integration_tests/model_tests.py @@ -36,7 +36,7 @@ from superset.common.db_query_status import QueryStatus from superset.models.core import Database from superset.models.slice import Slice from superset.models.sql_types.base import literal_dttm_type_factory -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from .base_tests import SupersetTestCase from .fixtures.energy_dashboard import ( diff --git a/tests/integration_tests/queries/api_tests.py b/tests/integration_tests/queries/api_tests.py index 45a807b7b2..d086e10821 100644 --- a/tests/integration_tests/queries/api_tests.py +++ b/tests/integration_tests/queries/api_tests.py @@ -29,7 +29,7 @@ import tests.integration_tests.test_app from superset import db, security_manager from superset.common.db_query_status import QueryStatus from superset.models.core import Database -from superset.utils.core import get_example_database, get_main_database +from superset.utils.database import get_example_database, get_main_database from superset.models.sql_lab import Query from tests.integration_tests.base_tests import SupersetTestCase diff --git a/tests/integration_tests/queries/saved_queries/api_tests.py b/tests/integration_tests/queries/saved_queries/api_tests.py index cf698d370a..2659bc224f 100644 --- a/tests/integration_tests/queries/saved_queries/api_tests.py +++ b/tests/integration_tests/queries/saved_queries/api_tests.py @@ -31,7 +31,7 @@ from superset import db from superset.models.core import Database from superset.models.core import FavStar from superset.models.sql_lab import SavedQuery -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.fixtures.importexport import ( diff --git a/tests/integration_tests/queries/saved_queries/commands_tests.py b/tests/integration_tests/queries/saved_queries/commands_tests.py index 1bf4b26cab..f90924ba0e 100644 --- a/tests/integration_tests/queries/saved_queries/commands_tests.py +++ b/tests/integration_tests/queries/saved_queries/commands_tests.py @@ -30,7 +30,7 @@ from superset.queries.saved_queries.commands.export import ExportSavedQueriesCom from superset.queries.saved_queries.commands.importers.v1 import ( ImportSavedQueriesCommand, ) -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.fixtures.importexport import ( database_config, diff --git a/tests/integration_tests/reports/api_tests.py b/tests/integration_tests/reports/api_tests.py index f7caceffa0..af01556f9f 100644 --- a/tests/integration_tests/reports/api_tests.py +++ b/tests/integration_tests/reports/api_tests.py @@ -38,7 +38,7 @@ from superset.models.reports import ( ReportRecipientType, ReportState, ) -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.conftest import with_feature_flags from tests.integration_tests.fixtures.birth_names_dashboard import ( diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py index 2c963f9351..d91d2cd158 100644 --- a/tests/integration_tests/reports/commands_tests.py +++ b/tests/integration_tests/reports/commands_tests.py @@ -56,7 +56,7 @@ from superset.reports.commands.exceptions import ( ) from superset.reports.commands.execute import AsyncExecuteReportScheduleCommand from superset.reports.commands.log_prune import AsyncPruneReportScheduleLogCommand -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, load_birth_names_data, diff --git a/tests/integration_tests/security_tests.py b/tests/integration_tests/security_tests.py index 9863e1b316..ab4450b04c 100644 --- a/tests/integration_tests/security_tests.py +++ b/tests/integration_tests/security_tests.py @@ -40,9 +40,9 @@ from superset.models.slice import Slice from superset.sql_parse import Table from superset.utils.core import ( backend, - get_example_database, get_example_default_schema, ) +from superset.utils.database import get_example_database from superset.views.access_requests import AccessRequestsModelView from .base_tests import SupersetTestCase diff --git a/tests/integration_tests/sql_validator_tests.py b/tests/integration_tests/sql_validator_tests.py index 92d5cf4061..3468b8c047 100644 --- a/tests/integration_tests/sql_validator_tests.py +++ b/tests/integration_tests/sql_validator_tests.py @@ -30,7 +30,7 @@ from superset.sql_validators.presto_db import ( PrestoDBSQLValidator, PrestoSQLValidationError, ) -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from .base_tests import SupersetTestCase diff --git a/tests/integration_tests/sqla_models_tests.py b/tests/integration_tests/sqla_models_tests.py index 00d3859009..eb54507800 100644 --- a/tests/integration_tests/sqla_models_tests.py +++ b/tests/integration_tests/sqla_models_tests.py @@ -39,9 +39,9 @@ from superset.utils.core import ( AdhocMetricExpressionType, FilterOperator, GenericDataType, - get_example_database, TemporalType, ) +from superset.utils.database import get_example_database from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, load_birth_names_data, diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py index 3c899d7d31..c96b7449b9 100644 --- a/tests/integration_tests/sqllab_tests.py +++ b/tests/integration_tests/sqllab_tests.py @@ -50,9 +50,8 @@ from superset.sql_parse import CtasMethod from superset.utils.core import ( backend, datetime_to_epoch, - get_example_database, - get_main_database, ) +from superset.utils.database import get_example_database, get_main_database from .base_tests import SupersetTestCase from .conftest import CTAS_SCHEMA_NAME diff --git a/tests/integration_tests/strategy_tests.py b/tests/integration_tests/strategy_tests.py index 59430f80b8..aec73b1efe 100644 --- a/tests/integration_tests/strategy_tests.py +++ b/tests/integration_tests/strategy_tests.py @@ -30,7 +30,7 @@ import pytest import pandas as pd from superset.models.slice import Slice -from superset.utils.core import get_example_database +from superset.utils.database import get_example_database from superset import db diff --git a/tests/integration_tests/utils_tests.py b/tests/integration_tests/utils_tests.py index fc9dba4c07..4bbb1a474f 100644 --- a/tests/integration_tests/utils_tests.py +++ b/tests/integration_tests/utils_tests.py @@ -54,7 +54,6 @@ from superset.utils.core import ( get_form_data_token, get_iterable, get_email_address_list, - get_or_create_db, get_stacktrace, json_int_dttm_ser, json_iso_dttm_ser, @@ -72,6 +71,7 @@ from superset.utils.core import ( zlib_compress, zlib_decompress, ) +from superset.utils.database import get_or_create_db from superset.utils import schema from superset.utils.hashing import md5_sha_from_str from superset.views.utils import (