From 4b8c12a87f8b1a80d54258d56523251cc7f6bd2c Mon Sep 17 00:00:00 2001
From: Gnought <1684105+gnought@users.noreply.github.com>
Date: Thu, 9 Nov 2023 03:56:38 +0800
Subject: [PATCH] chore: remove deprecated functions in SQLAlchemy (#25911)

---
 superset/db_engine_specs/__init__.py          | 39 +++++++++----------
 ...l_column_data_type_in_query_mysql_table.py |  2 +-
 ...hange_table_schema_description_to_long_.py |  2 +-
 3 files changed, 20 insertions(+), 23 deletions(-)

diff --git a/superset/db_engine_specs/__init__.py b/superset/db_engine_specs/__init__.py
index d4ec199133d38..9ec0f0416bd09 100644
--- a/superset/db_engine_specs/__init__.py
+++ b/superset/db_engine_specs/__init__.py
@@ -35,11 +35,10 @@
 from pathlib import Path
 from typing import Any, Optional
 
-import sqlalchemy.databases
 import sqlalchemy.dialects
 from importlib_metadata import entry_points
 from sqlalchemy.engine.default import DefaultDialect
-from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import NoSuchModuleError
 
 from superset import app, feature_flag_manager
 from superset.db_engine_specs.base import BaseEngineSpec
@@ -128,27 +127,26 @@ def get_available_engine_specs() -> dict[type[BaseEngineSpec], set[str]]:
     drivers: dict[str, set[str]] = defaultdict(set)
 
     # native SQLAlchemy dialects
-    for attr in sqlalchemy.databases.__all__:
-        dialect = getattr(sqlalchemy.dialects, attr)
-        for attribute in dialect.__dict__.values():
+    for attr in sqlalchemy.dialects.__all__:
+        try:
+            dialect = sqlalchemy.dialects.registry.load(attr)
             if (
-                hasattr(attribute, "dialect")
-                and inspect.isclass(attribute.dialect)
-                and issubclass(attribute.dialect, DefaultDialect)
+                issubclass(dialect, DefaultDialect)
+                and hasattr(dialect, "driver")
                 # adodbapi dialect is removed in SQLA 1.4 and doesn't implement the
                 # `dbapi` method, hence needs to be ignored to avoid logging a warning
-                and attribute.dialect.driver != "adodbapi"
+                and dialect.driver != "adodbapi"
             ):
                 try:
-                    attribute.dialect.dbapi()
+                    dialect.dbapi()
                 except ModuleNotFoundError:
                     continue
                 except Exception as ex:  # pylint: disable=broad-except
-                    logger.warning(
-                        "Unable to load dialect %s: %s", attribute.dialect, ex
-                    )
+                    logger.warning("Unable to load dialect %s: %s", dialect, ex)
                     continue
-                drivers[attr].add(attribute.dialect.driver)
+                drivers[attr].add(dialect.driver)
+        except NoSuchModuleError:
+            continue
 
     # installed 3rd-party dialects
     for ep in entry_points(group="sqlalchemy.dialects"):
@@ -167,21 +165,20 @@ def get_available_engine_specs() -> dict[type[BaseEngineSpec], set[str]]:
                 driver = driver.decode()
             drivers[backend].add(driver)
 
+    dbs_denylist = app.config["DBS_AVAILABLE_DENYLIST"]
+    if not feature_flag_manager.is_feature_enabled("ENABLE_SUPERSET_META_DB"):
+        dbs_denylist["superset"] = {""}
+    dbs_denylist_engines = dbs_denylist.keys()
     available_engines = {}
+
     for engine_spec in load_engine_specs():
         driver = drivers[engine_spec.engine]
-
-        # do not add denied db engine specs to available list
-        dbs_denylist = app.config["DBS_AVAILABLE_DENYLIST"]
-        if not feature_flag_manager.is_feature_enabled("ENABLE_SUPERSET_META_DB"):
-            dbs_denylist["superset"] = {""}
-        dbs_denylist_engines = dbs_denylist.keys()
-
         if (
             engine_spec.engine in dbs_denylist_engines
             and hasattr(engine_spec, "default_driver")
             and engine_spec.default_driver in dbs_denylist[engine_spec.engine]
         ):
+            # do not add denied db engine specs to available list
             continue
 
         # lookup driver by engine aliases.
diff --git a/superset/migrations/versions/2019-05-06_14-30_afc69274c25a_alter_sql_column_data_type_in_query_mysql_table.py b/superset/migrations/versions/2019-05-06_14-30_afc69274c25a_alter_sql_column_data_type_in_query_mysql_table.py
index 5d3a76798df81..9fb9acae7a90e 100644
--- a/superset/migrations/versions/2019-05-06_14-30_afc69274c25a_alter_sql_column_data_type_in_query_mysql_table.py
+++ b/superset/migrations/versions/2019-05-06_14-30_afc69274c25a_alter_sql_column_data_type_in_query_mysql_table.py
@@ -24,7 +24,7 @@
 """
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy.databases import mysql
+from sqlalchemy.dialects import mysql
 from sqlalchemy.dialects.mysql.base import MySQLDialect
 
 # revision identifiers, used by Alembic.
diff --git a/superset/migrations/versions/2019-12-03_13-50_89115a40e8ea_change_table_schema_description_to_long_.py b/superset/migrations/versions/2019-12-03_13-50_89115a40e8ea_change_table_schema_description_to_long_.py
index a1d37ab40f331..7679baa5aa337 100644
--- a/superset/migrations/versions/2019-12-03_13-50_89115a40e8ea_change_table_schema_description_to_long_.py
+++ b/superset/migrations/versions/2019-12-03_13-50_89115a40e8ea_change_table_schema_description_to_long_.py
@@ -28,7 +28,7 @@
 
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy.databases import mysql
+from sqlalchemy.dialects import mysql
 from sqlalchemy.dialects.mysql.base import MySQLDialect