Moving get_main_db out of tests and into superset utils (#5864)

superset/cli.py

@@ -20,7 +20,6 @@ import yaml
 from superset import (
     app, data, db, dict_import_export_util, security_manager, utils,
 )
-from tests.utils import get_main_database
 
 config = app.config
 celery_app = utils.get_celery_app(config)
@@ -390,7 +389,7 @@ def load_test_users_run():
         for perm in security_manager.find_role('Gamma').permissions:
             security_manager.add_permission_role(gamma_sqllab_role, perm)
         utils.get_or_create_main_db()
-        db_perm = get_main_database(security_manager.get_session).perm
+        db_perm = utils.get_main_database(security_manager.get_session).perm
         security_manager.merge_perm('database_access', db_perm)
         db_pvm = security_manager.find_permission_view_menu(
             view_menu_name=db_perm, permission_name='database_access')

superset/utils.py

@@ -836,10 +836,7 @@ def get_or_create_main_db():
     from superset.models import core as models
 
     logging.info('Creating database reference')
-    dbobj = (
+    dbobj = get_main_database(db.session)
-        db.session.query(models.Database)
-        .filter_by(database_name='main')
-        .first())
     if not dbobj:
         dbobj = models.Database(database_name='main')
     dbobj.set_sqlalchemy_uri(conf.get('SQLALCHEMY_DATABASE_URI'))
@@ -850,6 +847,15 @@ def get_or_create_main_db():
     return dbobj
 
 
+def get_main_database(session):
+    from superset.models import core as models
+    return (
+        session.query(models.Database)
+        .filter_by(database_name='main')
+        .first()
+    )
+
+
 def is_adhoc_metric(metric):
     return (
         isinstance(metric, dict) and

tests/base_tests.py

@@ -18,7 +18,7 @@ from superset import app, cli, db, security_manager
 from superset.connectors.druid.models import DruidCluster, DruidDatasource
 from superset.connectors.sqla.models import SqlaTable
 from superset.models import core as models
-from .utils import get_main_database
+from superset.utils import get_main_database
 
 
 BASE_DIR = app.config.get('BASE_DIR')

tests/celery_tests.py

@@ -18,8 +18,8 @@ from superset import app, cli, db, security_manager
 from superset.models.helpers import QueryStatus
 from superset.models.sql_lab import Query
 from superset.sql_parse import SupersetQuery
+from superset.utils import get_main_database
 from .base_tests import SupersetTestCase
-from .utils import get_main_database
 
 
 BASE_DIR = app.config.get('BASE_DIR')

tests/core_tests.py

@@ -27,9 +27,9 @@ from superset.connectors.sqla.models import SqlaTable
 from superset.db_engine_specs import BaseEngineSpec
 from superset.models import core as models
 from superset.models.sql_lab import Query
+from superset.utils import get_main_database
 from superset.views.core import DatabaseView
 from .base_tests import SupersetTestCase
-from .utils import get_main_database
 
 
 class CoreTests(SupersetTestCase):

tests/dict_import_export_tests.py

@@ -15,8 +15,8 @@ from superset.connectors.druid.models import (
     DruidColumn, DruidDatasource, DruidMetric,
 )
 from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn
+from superset.utils import get_main_database
 from .base_tests import SupersetTestCase
-from .utils import get_main_database
 
 DBREF = 'dict_import__export_test'
 NAME_PREFIX = 'dict_'

tests/model_tests.py

@@ -10,8 +10,8 @@ from sqlalchemy.engine.url import make_url
 
 from superset import app, db
 from superset.models.core import Database
+from superset.utils import get_main_database
 from .base_tests import SupersetTestCase
-from .utils import get_main_database
 
 
 class DatabaseModelTestCase(SupersetTestCase):

tests/sqllab_tests.py

@@ -15,8 +15,8 @@ from superset import db, security_manager, utils
 from superset.dataframe import SupersetDataFrame
 from superset.db_engine_specs import BaseEngineSpec
 from superset.models.sql_lab import Query
+from superset.utils import get_main_database
 from .base_tests import SupersetTestCase
-from .utils import get_main_database
 
 
 class SqlLabTests(SupersetTestCase):

tests/utils.py

@@ -7,20 +7,9 @@ from __future__ import unicode_literals
 import json
 from os import path
 
-from superset import db
-from superset.models import core as models
-
 FIXTURES_DIR = 'tests/fixtures'
 
 
 def load_fixture(fixture_file_name):
     with open(path.join(FIXTURES_DIR, fixture_file_name)) as fixture_file:
         return json.load(fixture_file)
-
-
-def get_main_database(session):
-    return (
-        db.session.query(models.Database)
-        .filter_by(database_name='main')
-        .first()
-    )