chore(supeset.utils.core): move all database utils to database utils …

…module (#18058)

superset/cli.py

@@ -37,9 +37,9 @@
 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)
 

superset/commands/importers/v1/examples.py

@@ -42,7 +42,8 @@
 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):

superset/examples/bart_lines.py

@@ -21,8 +21,8 @@
 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
 
 

superset/examples/birth_names.py

@@ -29,8 +29,8 @@
 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,

superset/examples/country_map.py

@@ -20,10 +20,10 @@
 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 @@
 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)

superset/examples/energy.py

@@ -21,10 +21,10 @@
 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)

superset/examples/flights.py

@@ -17,16 +17,16 @@
 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
 
 
 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)

superset/examples/long_lat.py

@@ -21,9 +21,9 @@
 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 @@
 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)

superset/examples/multiformat_time_series.py

@@ -21,8 +21,8 @@
 
 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,

superset/examples/paris.py

@@ -19,15 +19,15 @@
 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)

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)

superset/examples/sf_population_polygons.py

@@ -19,8 +19,8 @@
 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)

superset/examples/world_bank.py

@@ -23,6 +23,7 @@
 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)

superset/utils/core.py

@@ -86,7 +86,6 @@
 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 @@
     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 @@
 
 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.

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)

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)