john-bodley commented on code in PR #19914:
URL: https://github.com/apache/superset/pull/19914#discussion_r862531806


##########
superset/databases/commands/test_connection.py:
##########
@@ -74,42 +75,43 @@ def run(self) -> None:
 
             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):

Review Comment:
   Same logic as previously, though indented due to the context manager.



##########
superset/db_engine_specs/base.py:
##########
@@ -393,10 +393,7 @@ def get_engine(
         schema: Optional[str] = None,
         source: Optional[str] = None,
     ) -> Engine:
-        user_name = utils.get_username()
-        return database.get_sqla_engine(
-            schema=schema, nullpool=True, user_name=user_name, source=source
-        )
+        return database.get_sqla_engine(schema=schema, source=source)

Review Comment:
   `nullpool=True` is the default.



##########
superset/databases/commands/validate.py:
##########
@@ -115,22 +116,23 @@ def run(self) -> None:
         )
         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):

Review Comment:
   Same logic as previously, though indented due to the context manager.



##########
superset/reports/commands/alert.py:
##########
@@ -145,18 +147,21 @@ def _execute_query(self) -> pd.DataFrame:
             limited_rendered_sql = 
self._report_schedule.database.apply_limit_to_sql(
                 rendered_sql, ALERT_SQL_LIMIT
             )
-            query_username = app.config["THUMBNAIL_SELENIUM_USER"]
-            start = default_timer()
-            df = self._report_schedule.database.get_df(
-                sql=limited_rendered_sql, username=query_username
-            )
-            stop = default_timer()
-            logger.info(
-                "Query for %s took %.2f ms",
-                self._report_schedule.name,
-                (stop - start) * 1000.0,
-            )
-            return df
+
+            with override_user(

Review Comment:
   Same logic as previously, though indented due to the context manager.



##########
superset/sql_lab.py:
##########
@@ -149,37 +156,35 @@ def get_sql_results(  # pylint: disable=too-many-arguments
     rendered_query: str,
     return_results: bool = True,
     store_results: bool = False,
-    user_name: Optional[str] = None,
+    username: Optional[str] = None,
     start_time: Optional[float] = None,
     expand_data: bool = False,
     log_params: Optional[Dict[str, Any]] = None,
 ) -> Optional[Dict[str, Any]]:
     """Executes the sql query returns the results."""
     with session_scope(not ctask.request.called_directly) as session:
-
-        try:
-            return execute_sql_statements(
-                query_id,
-                rendered_query,
-                return_results,
-                store_results,
-                user_name,
-                session=session,
-                start_time=start_time,
-                expand_data=expand_data,
-                log_params=log_params,
-            )
-        except Exception as ex:  # pylint: disable=broad-except
-            logger.debug("Query %d: %s", query_id, ex)
-            stats_logger.incr("error_sqllab_unhandled")
-            query = get_query(query_id, session)
-            return handle_query_error(ex, query, session)
+        with 
override_user(current_app.app_builder.sm.find_user(username=username)):

Review Comment:
   Same logic as previously, though indented due to the context manager.



##########
superset/sql_lab.py:
##########
@@ -149,37 +156,35 @@ def get_sql_results(  # pylint: disable=too-many-arguments
     rendered_query: str,
     return_results: bool = True,
     store_results: bool = False,
-    user_name: Optional[str] = None,
+    username: Optional[str] = None,

Review Comment:
   We still need to pass in the `username` to the async Celery task and then 
override.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to