diff --git a/superset/charts/api.py b/superset/charts/api.py index 724e981794..e3f3936f93 100644 --- a/superset/charts/api.py +++ b/superset/charts/api.py @@ -595,7 +595,7 @@ class ChartRestApi(BaseSupersetModelRestApi): except AsyncQueryTokenException: return self.response_401() - result = command.run_async() + result = command.run_async(g.user.get_id()) return self.response(202, **result) return self.get_data_response(command) diff --git a/superset/charts/commands/data.py b/superset/charts/commands/data.py index 275a723d60..1a383827cc 100644 --- a/superset/charts/commands/data.py +++ b/superset/charts/commands/data.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. import logging -from typing import Any, Dict +from typing import Any, Dict, Optional from flask import Request from marshmallow import ValidationError @@ -67,8 +67,8 @@ class ChartDataCommand(BaseCommand): return return_value - def run_async(self) -> Dict[str, Any]: - job_metadata = async_query_manager.init_job(self._async_channel_id) + def run_async(self, user_id: Optional[str]) -> Dict[str, Any]: + job_metadata = async_query_manager.init_job(self._async_channel_id, user_id) load_chart_data_into_cache.delay(job_metadata, self._form_data) return job_metadata diff --git a/superset/utils/async_query_manager.py b/superset/utils/async_query_manager.py index 77373b6857..90428a11ee 100644 --- a/superset/utils/async_query_manager.py +++ b/superset/utils/async_query_manager.py @@ -21,7 +21,7 @@ from typing import Any, Dict, List, Optional, Tuple import jwt import redis -from flask import Flask, request, Request, Response, session +from flask import Flask, g, request, Request, Response, session logger = logging.getLogger(__name__) @@ -34,11 +34,13 @@ class AsyncQueryJobException(Exception): pass -def build_job_metadata(channel_id: str, job_id: str, **kwargs: Any) -> Dict[str, Any]: +def build_job_metadata( + channel_id: str, job_id: str, user_id: Optional[str], **kwargs: Any +) -> Dict[str, Any]: return { "channel_id": channel_id, "job_id": job_id, - "user_id": session.get("user_id"), + "user_id": int(user_id) if user_id else None, "status": kwargs.get("status"), "errors": kwargs.get("errors", []), "result_url": kwargs.get("result_url"), @@ -115,7 +117,13 @@ class AsyncQueryManager: def validate_session( # pylint: disable=unused-variable response: Response, ) -> Response: - user_id = session["user_id"] if "user_id" in session else None + user_id = None + + try: + user_id = g.user.get_id() + user_id = int(user_id) + except Exception: # pylint: disable=broad-except + pass reset_token = ( not request.cookies.get(self._jwt_cookie_name) @@ -161,9 +169,11 @@ class AsyncQueryManager: logger.warning(exc) raise AsyncQueryTokenException("Failed to parse token") - def init_job(self, channel_id: str) -> Dict[str, Any]: + def init_job(self, channel_id: str, user_id: Optional[str]) -> Dict[str, Any]: job_id = str(uuid.uuid4()) - return build_job_metadata(channel_id, job_id, status=self.STATUS_PENDING) + return build_job_metadata( + channel_id, job_id, user_id, status=self.STATUS_PENDING + ) def read_events( self, channel: str, last_id: Optional[str] diff --git a/superset/views/core.py b/superset/views/core.py index e1b98e6db4..622b1170a7 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -611,7 +611,9 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods async_channel_id = async_query_manager.parse_jwt_from_request( request )["channel"] - job_metadata = async_query_manager.init_job(async_channel_id) + job_metadata = async_query_manager.init_job( + async_channel_id, g.user.get_id() + ) load_explore_json_into_cache.delay( job_metadata, form_data, response_type, force )