Fix sorting of DBs

apache · May 25, 2021 · dbab3dd · dbab3dd
1 parent e546520
commit dbab3dd
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 25 deletions.
diff --git a/superset/config.py b/superset/config.py
@@ -1073,13 +1073,13 @@ def CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC(
 
 # A list of preferred databases, in order. These databases will be
 # displayed prominently in the "Add Database" dialog. You should
-# use the "engine" attribute of the corresponding DB engine spec in
-# `superset/db_engine_specs/`.
+# use the "engine_name" attribute of the corresponding DB engine spec
+# in `superset/db_engine_specs/`.
 PREFERRED_DATABASES: List[str] = [
-    # "postgresql",
-    # "presto",
-    # "mysql",
-    # "sqlite",
+    # "PostgreSQL",
+    # "Presto",
+    # "MySQL",
+    # "SQLite",
     # etc.
 ]
 

diff --git a/superset/databases/api.py b/superset/databases/api.py
@@ -917,7 +917,7 @@ def available(self) -> Response:
             payload: Dict[str, Any] = {
                 "name": engine_spec.engine_name,
                 "engine": engine_spec.engine,
-                "available_drivers": list(drivers),
+                "available_drivers": sorted(drivers),
                 "preferred": engine_spec.engine_name in preferred_databases,
             }
 
@@ -939,13 +939,25 @@ def available(self) -> Response:
 
             available_databases.append(payload)
 
-        available_databases.sort(
-            key=lambda payload: preferred_databases.index(payload["engine"])
-            if payload["engine"] in preferred_databases
-            else len(preferred_databases)
+        # sort preferred first
+        response = sorted(
+            (payload for payload in available_databases if payload["preferred"]),
+            key=lambda payload: preferred_databases.index(payload["name"]),
         )
 
-        return self.response(200, databases=available_databases)
+        # add others
+        response.extend(
+            sorted(
+                (
+                    payload
+                    for payload in available_databases
+                    if not payload["preferred"]
+                ),
+                key=lambda payload: payload["name"],
+            )
+        )
+
+        return self.response(200, databases=response)
 
     @expose("/validate_parameters", methods=["POST"])
     @protect()

diff --git a/superset/db_engine_specs/__init__.py b/superset/db_engine_specs/__init__.py
@@ -107,7 +107,7 @@ def get_available_engine_specs() -> Dict[Type[BaseEngineSpec], Set[str]]:
                     attribute.dialect.dbapi()
                 except ModuleNotFoundError:
                     continue
-                except Exception as ex:
+                except Exception as ex:  # pylint: disable=broad-except
                     logger.warning(
                         "Unable to load dialect %s: %s", attribute.dialect, ex
                     )

diff --git a/tests/databases/api_tests.py b/tests/databases/api_tests.py
@@ -1372,7 +1372,7 @@ def test_function_names(self, mock_get_function_names):
     @mock.patch("superset.databases.api.get_available_engine_specs")
     @mock.patch("superset.databases.api.app")
     def test_available(self, app, get_available_engine_specs):
-        app.config = {"PREFERRED_DATABASES": ["postgresql"]}
+        app.config = {"PREFERRED_DATABASES": ["PostgreSQL", "Google BigQuery"]}
         get_available_engine_specs.return_value = {
             PostgresEngineSpec: {"psycopg2"},
             BigQueryEngineSpec: {"bigquery"},
@@ -1432,7 +1432,7 @@ def test_available(self, app, get_available_engine_specs):
                         "required": ["database", "host", "port", "username"],
                         "type": "object",
                     },
-                    "preferred": False,
+                    "preferred": True,
                     "sqlalchemy_uri_placeholder": "postgresql://user:password@host:port/dbname[?key=value&key=value...]",
                 },
                 {
@@ -1450,11 +1450,18 @@ def test_available(self, app, get_available_engine_specs):
                         },
                         "type": "object",
                     },
-                    "preferred": False,
+                    "preferred": True,
                     "sqlalchemy_uri_placeholder": "bigquery://{project_id}",
                 },
                 {
-                    "available_drivers": ["mysqldb", "mysqlconnector"],
+                    "available_drivers": ["psycopg2"],
+                    "default_driver": "",
+                    "engine": "redshift",
+                    "name": "Amazon Redshift",
+                    "preferred": False,
+                },
+                {
+                    "available_drivers": ["mysqlconnector", "mysqldb"],
                     "default_driver": "mysqldb",
                     "engine": "mysql",
                     "name": "MySQL",
@@ -1499,13 +1506,6 @@ def test_available(self, app, get_available_engine_specs):
                     "preferred": False,
                     "sqlalchemy_uri_placeholder": "mysql://user:password@host:port/dbname[?key=value&key=value...]",
                 },
-                {
-                    "available_drivers": ["psycopg2"],
-                    "default_driver": "",
-                    "engine": "redshift",
-                    "name": "Amazon Redshift",
-                    "preferred": False,
-                },
                 {
                     "available_drivers": [""],
                     "engine": "hana",
@@ -1518,7 +1518,7 @@ def test_available(self, app, get_available_engine_specs):
     @mock.patch("superset.databases.api.get_available_engine_specs")
     @mock.patch("superset.databases.api.app")
     def test_available_no_default(self, app, get_available_engine_specs):
-        app.config = {"PREFERRED_DATABASES": ["mysql"]}
+        app.config = {"PREFERRED_DATABASES": ["MySQL"]}
         get_available_engine_specs.return_value = {
             MySQLEngineSpec: {"mysqlconnector"},
             HanaEngineSpec: {""},