chore: upgrade SQLAlchemy to 1.4 (#19890)

* chore: upgrade SQLAlchemy

* Convert integration test to unit test

* Fix SQLite

* Update method names/docstrings

* Skip test

* Fix SQLite

requirements/base.txt

@@ -122,6 +122,8 @@ geopy==2.2.0
     # via apache-superset
 graphlib-backport==1.0.3
     # via apache-superset
+greenlet==1.1.2
+    # via sqlalchemy
 gunicorn==20.1.0
     # via apache-superset
 hashids==1.3.1
@@ -259,7 +261,7 @@ six==1.16.0
     #   wtforms-json
 slackclient==2.5.0
     # via apache-superset
-sqlalchemy==1.3.24
+sqlalchemy==1.4.36
     # via
     #   alembic
     #   apache-superset

requirements/docker.txt

@@ -12,8 +12,6 @@
     #   -r requirements/docker.in
 gevent==21.8.0
     # via -r requirements/docker.in
-greenlet==1.1.1
-    # via gevent
 psycopg2-binary==2.9.1
     # via apache-superset
 zope-event==4.5.0

setup.py

@@ -109,7 +109,7 @@ def get_git_sha() -> str:
         "selenium>=3.141.0",
         "simplejson>=3.15.0",
         "slackclient==2.5.0",  # PINNED! slack changes file upload api in the future versions
-        "sqlalchemy>=1.3.16, <1.4, !=1.3.21",
+        "sqlalchemy>=1.4, <2",
         "sqlalchemy-utils>=0.37.8, <0.38",
         "sqlparse==0.3.0",  # PINNED! see https://github.com/andialbrecht/sqlparse/issues/562
         "tabulate==0.8.9",

superset/commands/importers/v1/assets.py

@@ -129,7 +129,6 @@ def _import(session: Session, configs: Dict[str, Any]) -> None:
             {"dashboard_id": dashboard_id, "slice_id": chart_id}
             for (dashboard_id, chart_id) in dashboard_chart_ids
         ]
-        # pylint: disable=no-value-for-parameter # sqlalchemy/issues/4656
         session.execute(dashboard_slices.insert(), values)
 
     def run(self) -> None:

superset/commands/importers/v1/examples.py

@@ -181,5 +181,4 @@ def _import(  # pylint: disable=arguments-differ, too-many-locals, too-many-bran
             {"dashboard_id": dashboard_id, "slice_id": chart_id}
             for (dashboard_id, chart_id) in dashboard_chart_ids
         ]
-        # pylint: disable=no-value-for-parameter # sqlalchemy/issues/4656
         session.execute(dashboard_slices.insert(), values)

superset/dashboards/commands/importers/v1/__init__.py

@@ -139,5 +139,4 @@ def _import(
             {"dashboard_id": dashboard_id, "slice_id": chart_id}
             for (dashboard_id, chart_id) in dashboard_chart_ids
         ]
-        # pylint: disable=no-value-for-parameter # sqlalchemy/issues/4656
         session.execute(dashboard_slices.insert(), values)

superset/db_engine_specs/base.py

@@ -117,7 +117,9 @@ class TimeGrain(NamedTuple):
 }
 
 
-class TimestampExpression(ColumnClause):  # pylint: disable=abstract-method
+class TimestampExpression(
+    ColumnClause
+):  # pylint: disable=abstract-method, too-many-ancestors
     def __init__(self, expr: str, col: ColumnClause, **kwargs: Any) -> None:
         """Sqlalchemy class that can be can be used to render native column elements
         respeting engine-specific quoting rules as part of a string-based expression.
@@ -933,9 +935,13 @@ def extract_errors(
         ]
 
     @classmethod
-    def adjust_database_uri(cls, uri: URL, selected_schema: Optional[str]) -> None:
+    def adjust_database_uri(  # pylint: disable=unused-argument
+        cls,
+        uri: URL,
+        selected_schema: Optional[str],
+    ) -> URL:
         """
-        Mutate the database component of the SQLAlchemy URI.
+        Return a modified URL with a new database component.
 
         The URI here represents the URI as entered when saving the database,
         ``selected_schema`` is the schema currently active presumably in
@@ -949,9 +955,10 @@ def adjust_database_uri(cls, uri: URL, selected_schema: Optional[str]) -> None:
         For those it's probably better to not alter the database
         component of the URI with the schema name, it won't work.
 
-        Some database drivers like presto accept '{catalog}/{schema}' in
+        Some database drivers like Presto accept '{catalog}/{schema}' in
         the database component of the URL, that can be handled here.
         """
+        return uri
 
     @classmethod
     def patch(cls) -> None:
@@ -1206,17 +1213,20 @@ def estimate_query_cost(
         return costs
 
     @classmethod
-    def modify_url_for_impersonation(
+    def get_url_for_impersonation(
         cls, url: URL, impersonate_user: bool, username: Optional[str]
-    ) -> None:
+    ) -> URL:
         """
-        Modify the SQL Alchemy URL object with the user to impersonate if applicable.
+        Return a modified URL with the username set.
+
         :param url: SQLAlchemy URL object
         :param impersonate_user: Flag indicating if impersonation is enabled
         :param username: Effective username
         """
         if impersonate_user and username is not None:
-            url.username = username
+            url = url.set(username=username)
+
+        return url
 
     @classmethod
     def update_impersonation_config(

superset/db_engine_specs/drill.py

@@ -68,26 +68,31 @@ def convert_dttm(
         return None
 
     @classmethod
-    def adjust_database_uri(cls, uri: URL, selected_schema: Optional[str]) -> None:
+    def adjust_database_uri(cls, uri: URL, selected_schema: Optional[str]) -> URL:
         if selected_schema:
-            uri.database = parse.quote(selected_schema, safe="")
+            uri = uri.set(database=parse.quote(selected_schema, safe=""))
+
+        return uri
 
     @classmethod
-    def modify_url_for_impersonation(
+    def get_url_for_impersonation(
         cls, url: URL, impersonate_user: bool, username: Optional[str]
-    ) -> None:
+    ) -> URL:
         """
-        Modify the SQL Alchemy URL object with the user to impersonate if applicable.
+        Return a modified URL with the username set.
+
         :param url: SQLAlchemy URL object
         :param impersonate_user: Flag indicating if impersonation is enabled
         :param username: Effective username
         """
         if impersonate_user and username is not None:
             if url.drivername == "drill+odbc":
-                url.query["DelegationUID"] = username
+                url = url.update_query_dict({"DelegationUID": username})
             elif url.drivername in ["drill+sadrill", "drill+jdbc"]:
-                url.query["impersonation_target"] = username
+                url = url.update_query_dict({"impersonation_target": username})
             else:
                 raise SupersetDBAPIProgrammingError(
                     f"impersonation is not supported for {url.drivername}"
                 )
+
+        return url

superset/db_engine_specs/gsheets.py

@@ -81,16 +81,18 @@ class GSheetsEngineSpec(SqliteEngineSpec):
     }
 
     @classmethod
-    def modify_url_for_impersonation(
+    def get_url_for_impersonation(
         cls,
         url: URL,
         impersonate_user: bool,
         username: Optional[str],
-    ) -> None:
+    ) -> URL:
         if impersonate_user and username is not None:
             user = security_manager.find_user(username=username)
             if user and user.email:
-                url.query["subject"] = user.email
+                url = url.update_query_dict({"subject": user.email})
+
+        return url
 
     @classmethod
     def extra_table_metadata(

superset/db_engine_specs/hive.py

@@ -269,9 +269,11 @@ def epoch_to_dttm(cls) -> str:
     @classmethod
     def adjust_database_uri(
         cls, uri: URL, selected_schema: Optional[str] = None
-    ) -> None:
+    ) -> URL:
         if selected_schema:
-            uri.database = parse.quote(selected_schema, safe="")
+            uri = uri.set(database=parse.quote(selected_schema, safe=""))
+
+        return uri
 
     @classmethod
     def _extract_error_message(cls, ex: Exception) -> str:
@@ -485,17 +487,19 @@ def select_star(  # pylint: disable=too-many-arguments
         )
 
     @classmethod
-    def modify_url_for_impersonation(
+    def get_url_for_impersonation(
         cls, url: URL, impersonate_user: bool, username: Optional[str]
-    ) -> None:
+    ) -> URL:
         """
-        Modify the SQL Alchemy URL object with the user to impersonate if applicable.
+        Return a modified URL with the username set.
+
         :param url: SQLAlchemy URL object
         :param impersonate_user: Flag indicating if impersonation is enabled
         :param username: Effective username
         """
         # Do nothing in the URL object since instead this should modify
         # the configuraiton dictionary. See get_configuration_for_impersonation
+        return url
 
     @classmethod
     def update_impersonation_config(

superset/db_engine_specs/mysql.py

@@ -193,9 +193,11 @@ def convert_dttm(
     @classmethod
     def adjust_database_uri(
         cls, uri: URL, selected_schema: Optional[str] = None
-    ) -> None:
+    ) -> URL:
         if selected_schema:
-            uri.database = parse.quote(selected_schema, safe="")
+            uri = uri.set(database=parse.quote(selected_schema, safe=""))
+
+        return uri
 
     @classmethod
     def get_datatype(cls, type_code: Any) -> Optional[str]:

superset/db_engine_specs/presto.py

@@ -33,7 +33,7 @@
 from sqlalchemy import Column, literal_column, types
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.engine.reflection import Inspector
-from sqlalchemy.engine.result import RowProxy
+from sqlalchemy.engine.result import Row as ResultRow
 from sqlalchemy.engine.url import URL
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import ColumnClause, Select
@@ -430,7 +430,7 @@ def _parse_structural_column(  # pylint: disable=too-many-locals
     @classmethod
     def _show_columns(
         cls, inspector: Inspector, table_name: str, schema: Optional[str]
-    ) -> List[RowProxy]:
+    ) -> List[ResultRow]:
         """
         Show presto column names
         :param inspector: object that performs database schema inspection
@@ -729,15 +729,17 @@ def humanize(value: Any, suffix: str) -> str:
     @classmethod
     def adjust_database_uri(
         cls, uri: URL, selected_schema: Optional[str] = None
-    ) -> None:
+    ) -> URL:
         database = uri.database
         if selected_schema and database:
             selected_schema = parse.quote(selected_schema, safe="")
             if "/" in database:
                 database = database.split("/")[0] + "/" + selected_schema
             else:
                 database += "/" + selected_schema
-            uri.database = database
+            uri = uri.set(database=database)
+
+        return uri
 
     @classmethod
     def convert_dttm(

superset/db_engine_specs/snowflake.py

@@ -114,13 +114,15 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
     @classmethod
     def adjust_database_uri(
         cls, uri: URL, selected_schema: Optional[str] = None
-    ) -> None:
+    ) -> URL:
         database = uri.database
         if "/" in uri.database:
             database = uri.database.split("/")[0]
         if selected_schema:
             selected_schema = parse.quote(selected_schema, safe="")
-            uri.database = database + "/" + selected_schema
+            uri = uri.set(database=f"{database}/{selected_schema}")
+
+        return uri
 
     @classmethod
     def epoch_to_dttm(cls) -> str:

superset/db_engine_specs/trino.py

@@ -65,16 +65,18 @@ def update_impersonation_config(
             connect_args["user"] = username
 
     @classmethod
-    def modify_url_for_impersonation(
+    def get_url_for_impersonation(
         cls, url: URL, impersonate_user: bool, username: Optional[str]
-    ) -> None:
+    ) -> URL:
         """
-        Modify the SQL Alchemy URL object with the user to impersonate if applicable.
+        Return a modified URL with the username set.
+
         :param url: SQLAlchemy URL object
         :param impersonate_user: Flag indicating if impersonation is enabled
         :param username: Effective username
         """
         # Do nothing and let update_impersonation_config take care of impersonation
+        return url
 
     @classmethod
     def get_allow_cost_estimate(cls, extra: Dict[str, Any]) -> bool:

superset/migrations/versions/2021-08-31_11-37_021b81fe4fbb_add_type_to_native_filter_configuration.py

@@ -31,7 +31,7 @@
 
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy.ext.declarative.api import declarative_base
+from sqlalchemy.ext.declarative import declarative_base
 
 from superset import db
 

superset/models/core.py

@@ -312,15 +312,15 @@ def get_password_masked_url_from_uri(  # pylint: disable=invalid-name
     def get_password_masked_url(cls, masked_url: URL) -> URL:
         url_copy = deepcopy(masked_url)
         if url_copy.password is not None:
-            url_copy.password = PASSWORD_MASK
+            url_copy = url_copy.set(password=PASSWORD_MASK)
         return url_copy
 
     def set_sqlalchemy_uri(self, uri: str) -> None:
         conn = make_url_safe(uri.strip())
         if conn.password != PASSWORD_MASK and not custom_password_store:
             # do not over-write the password with the password mask
             self.password = conn.password
-        conn.password = PASSWORD_MASK if conn.password else None
+        conn = conn.set(password=PASSWORD_MASK if conn.password else None)
         self.sqlalchemy_uri = str(conn)  # hides the password
 
     def get_effective_user(self, object_url: URL) -> Optional[str]:
@@ -355,12 +355,12 @@ def get_sqla_engine(
     ) -> Engine:
         extra = self.get_extra()
         sqlalchemy_url = make_url_safe(self.sqlalchemy_uri_decrypted)
-        self.db_engine_spec.adjust_database_uri(sqlalchemy_url, schema)
+        sqlalchemy_url = self.db_engine_spec.adjust_database_uri(sqlalchemy_url, schema)
         effective_username = self.get_effective_user(sqlalchemy_url)
         # If using MySQL or Presto for example, will set url.username
         # If using Hive, will not do anything yet since that relies on a
         # configuration parameter instead.
-        self.db_engine_spec.modify_url_for_impersonation(
+        sqlalchemy_url = self.db_engine_spec.get_url_for_impersonation(
             sqlalchemy_url, self.impersonate_user, effective_username
         )
 
@@ -736,9 +736,9 @@ def sqlalchemy_uri_decrypted(self) -> str:
             # (so users see 500 less often)
             return "dialect://invalid_uri"
         if custom_password_store:
-            conn.password = custom_password_store(conn)
+            conn = conn.set(password=custom_password_store(conn))
         else:
-            conn.password = self.password
+            conn = conn.set(password=self.password)
         return str(conn)
 
     @property

superset/models/sql_types/presto_sql_types.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# pylint: disable=abstract-method
+# pylint: disable=abstract-method, no-init
 from typing import Any, Dict, List, Optional, Type
 
 from sqlalchemy.engine.interfaces import Dialect

superset/utils/encrypt.py

@@ -20,7 +20,7 @@
 
 from flask import Flask
 from sqlalchemy import text, TypeDecorator
-from sqlalchemy.engine import Connection, Dialect, RowProxy
+from sqlalchemy.engine import Connection, Dialect, Row
 from sqlalchemy_utils import EncryptedType
 
 logger = logging.getLogger(__name__)
@@ -114,13 +114,13 @@ def _read_bytes(col_name: str, value: Any) -> Optional[bytes]:
     @staticmethod
     def _select_columns_from_table(
         conn: Connection, column_names: List[str], table_name: str
-    ) -> RowProxy:
+    ) -> Row:
         return conn.execute(f"SELECT id, {','.join(column_names)} FROM {table_name}")
 
     def _re_encrypt_row(
         self,
         conn: Connection,
-        row: RowProxy,
+        row: Row,
         table_name: str,
         columns: Dict[str, EncryptedType],
     ) -> None: