fix: Refactor SQL username logic (#19914)

Co-authored-by: John Bodley <john.bodley@airbnb.com>

superset/databases/commands/test_connection.py

@@ -38,6 +38,7 @@ from superset.errors import ErrorLevel, SupersetErrorType
 from superset.exceptions import SupersetSecurityException, SupersetTimeoutException
 from superset.extensions import event_logger
 from superset.models.core import Database
+from superset.utils.core import override_user
 
 logger = logging.getLogger(__name__)
 
@@ -74,42 +75,43 @@ class TestConnectionDatabaseCommand(BaseCommand):
 
             database.set_sqlalchemy_uri(uri)
             database.db_engine_spec.mutate_db_for_connection_test(database)
-            username = self._actor.username if self._actor is not None else None
-            engine = database.get_sqla_engine(user_name=username)
-            event_logger.log_with_context(
-                action="test_connection_attempt",
-                engine=database.db_engine_spec.__name__,
-            )
-            with closing(engine.raw_connection()) as conn:
-                try:
-                    alive = func_timeout(
-                        int(
-                            app.config[
-                                "TEST_DATABASE_CONNECTION_TIMEOUT"
-                            ].total_seconds()
-                        ),
-                        engine.dialect.do_ping,
-                        args=(conn,),
-                    )
-                except (sqlite3.ProgrammingError, RuntimeError):
-                    # SQLite can't run on a separate thread, so ``func_timeout`` fails
-                    # RuntimeError catches the equivalent error from duckdb.
-                    alive = engine.dialect.do_ping(conn)
-                except FunctionTimedOut as ex:
-                    raise SupersetTimeoutException(
-                        error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
-                        message=(
-                            "Please check your connection details and database settings, "
-                            "and ensure that your database is accepting connections, "
-                            "then try connecting again."
-                        ),
-                        level=ErrorLevel.ERROR,
-                        extra={"sqlalchemy_uri": database.sqlalchemy_uri},
-                    ) from ex
-                except Exception:  # pylint: disable=broad-except
-                    alive = False
-                if not alive:
-                    raise DBAPIError(None, None, None)
+
+            with override_user(self._actor):
+                engine = database.get_sqla_engine()
+                event_logger.log_with_context(
+                    action="test_connection_attempt",
+                    engine=database.db_engine_spec.__name__,
+                )
+                with closing(engine.raw_connection()) as conn:
+                    try:
+                        alive = func_timeout(
+                            int(
+                                app.config[
+                                    "TEST_DATABASE_CONNECTION_TIMEOUT"
+                                ].total_seconds()
+                            ),
+                            engine.dialect.do_ping,
+                            args=(conn,),
+                        )
+                    except (sqlite3.ProgrammingError, RuntimeError):
+                        # SQLite can't run on a separate thread, so ``func_timeout`` fails
+                        # RuntimeError catches the equivalent error from duckdb.
+                        alive = engine.dialect.do_ping(conn)
+                    except FunctionTimedOut as ex:
+                        raise SupersetTimeoutException(
+                            error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
+                            message=(
+                                "Please check your connection details and database "
+                                "settings, and ensure that your database is accepting "
+                                "connections, then try connecting again."
+                            ),
+                            level=ErrorLevel.ERROR,
+                            extra={"sqlalchemy_uri": database.sqlalchemy_uri},
+                        ) from ex
+                    except Exception:  # pylint: disable=broad-except
+                        alive = False
+                    if not alive:
+                        raise DBAPIError(None, None, None)
 
             # Log succesful connection test with engine
             event_logger.log_with_context(

superset/databases/commands/validate.py

@@ -35,6 +35,7 @@ from superset.db_engine_specs.base import BasicParametersMixin
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.extensions import event_logger
 from superset.models.core import Database
+from superset.utils.core import override_user
 
 BYPASS_VALIDATION_ENGINES = {"bigquery"}
 
@@ -115,22 +116,23 @@ class ValidateDatabaseParametersCommand(BaseCommand):
         )
         database.set_sqlalchemy_uri(sqlalchemy_uri)
         database.db_engine_spec.mutate_db_for_connection_test(database)
-        username = self._actor.username if self._actor is not None else None
-        engine = database.get_sqla_engine(user_name=username)
-        try:
-            with closing(engine.raw_connection()) as conn:
-                alive = engine.dialect.do_ping(conn)
-        except Exception as ex:
-            url = make_url_safe(sqlalchemy_uri)
-            context = {
-                "hostname": url.host,
-                "password": url.password,
-                "port": url.port,
-                "username": url.username,
-                "database": url.database,
-            }
-            errors = database.db_engine_spec.extract_errors(ex, context)
-            raise DatabaseTestConnectionFailedError(errors) from ex
+
+        with override_user(self._actor):
+            engine = database.get_sqla_engine()
+            try:
+                with closing(engine.raw_connection()) as conn:
+                    alive = engine.dialect.do_ping(conn)
+            except Exception as ex:
+                url = make_url_safe(sqlalchemy_uri)
+                context = {
+                    "hostname": url.host,
+                    "password": url.password,
+                    "port": url.port,
+                    "username": url.username,
+                    "database": url.database,
+                }
+                errors = database.db_engine_spec.extract_errors(ex, context)
+                raise DatabaseTestConnectionFailedError(errors) from ex
 
         if not alive:
             raise DatabaseOfflineError(