| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| import logging |
| import uuid |
| from typing import Any, Literal, Optional |
| |
| import jwt |
| import redis |
| from flask import Flask, Request, request, Response, session |
| |
| from superset.utils import json |
| from superset.utils.core import get_user_id |
| |
| logger = logging.getLogger(__name__) |
| |
| |
| class AsyncQueryTokenException(Exception): |
| pass |
| |
| |
| class AsyncQueryJobException(Exception): |
| pass |
| |
| |
| def build_job_metadata( |
| channel_id: str, job_id: str, user_id: Optional[int], **kwargs: Any |
| ) -> dict[str, Any]: |
| return { |
| "channel_id": channel_id, |
| "job_id": job_id, |
| "user_id": user_id, |
| "status": kwargs.get("status"), |
| "errors": kwargs.get("errors", []), |
| "result_url": kwargs.get("result_url"), |
| } |
| |
| |
| def parse_event(event_data: tuple[str, dict[str, Any]]) -> dict[str, Any]: |
| event_id = event_data[0] |
| event_payload = event_data[1]["data"] |
| return {"id": event_id, **json.loads(event_payload)} |
| |
| |
| def increment_id(redis_id: str) -> str: |
| # redis stream IDs are in this format: '1607477697866-0' |
| try: |
| prefix, last = redis_id[:-1], int(redis_id[-1]) |
| return prefix + str(last + 1) |
| except Exception: # pylint: disable=broad-except |
| return redis_id |
| |
| |
| class AsyncQueryManager: |
| MAX_EVENT_COUNT = 100 |
| STATUS_PENDING = "pending" |
| STATUS_RUNNING = "running" |
| STATUS_ERROR = "error" |
| STATUS_DONE = "done" |
| |
| def __init__(self) -> None: |
| super().__init__() |
| self._redis: redis.Redis # type: ignore |
| self._stream_prefix: str = "" |
| self._stream_limit: Optional[int] |
| self._stream_limit_firehose: Optional[int] |
| self._jwt_cookie_name: str = "" |
| self._jwt_cookie_secure: bool = False |
| self._jwt_cookie_domain: Optional[str] |
| self._jwt_cookie_samesite: Optional[Literal["None", "Lax", "Strict"]] = None |
| self._jwt_secret: str |
| self._load_chart_data_into_cache_job: Any = None |
| # pylint: disable=invalid-name |
| self._load_explore_json_into_cache_job: Any = None |
| |
| def init_app(self, app: Flask) -> None: |
| config = app.config |
| if ( |
| config["CACHE_CONFIG"]["CACHE_TYPE"] == "null" |
| or config["DATA_CACHE_CONFIG"]["CACHE_TYPE"] == "null" |
| ): |
| raise Exception( # pylint: disable=broad-exception-raised |
| """ |
| Cache backends (CACHE_CONFIG, DATA_CACHE_CONFIG) must be configured |
| and non-null in order to enable async queries |
| """ |
| ) |
| |
| if len(config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"]) < 32: |
| raise AsyncQueryTokenException( |
| "Please provide a JWT secret at least 32 bytes long" |
| ) |
| |
| self._redis = redis.Redis( |
| **config["GLOBAL_ASYNC_QUERIES_REDIS_CONFIG"], decode_responses=True |
| ) |
| self._stream_prefix = config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] |
| self._stream_limit = config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_LIMIT"] |
| self._stream_limit_firehose = config[ |
| "GLOBAL_ASYNC_QUERIES_REDIS_STREAM_LIMIT_FIREHOSE" |
| ] |
| self._jwt_cookie_name = config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME"] |
| self._jwt_cookie_secure = config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_SECURE"] |
| self._jwt_cookie_samesite = config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_SAMESITE"] |
| self._jwt_cookie_domain = config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_DOMAIN"] |
| self._jwt_secret = config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"] |
| |
| if config["GLOBAL_ASYNC_QUERIES_REGISTER_REQUEST_HANDLERS"]: |
| self.register_request_handlers(app) |
| |
| # pylint: disable=import-outside-toplevel |
| from superset.tasks.async_queries import ( |
| load_chart_data_into_cache, |
| load_explore_json_into_cache, |
| ) |
| |
| self._load_chart_data_into_cache_job = load_chart_data_into_cache |
| self._load_explore_json_into_cache_job = load_explore_json_into_cache |
| |
| def register_request_handlers(self, app: Flask) -> None: |
| @app.after_request |
| def validate_session(response: Response) -> Response: |
| user_id = get_user_id() |
| |
| reset_token = ( |
| not request.cookies.get(self._jwt_cookie_name) |
| or "async_channel_id" not in session |
| or "async_user_id" not in session |
| or user_id != session["async_user_id"] |
| ) |
| |
| if reset_token: |
| async_channel_id = str(uuid.uuid4()) |
| session["async_channel_id"] = async_channel_id |
| session["async_user_id"] = user_id |
| |
| sub = str(user_id) if user_id else None |
| token = jwt.encode( |
| {"channel": async_channel_id, "sub": sub}, |
| self._jwt_secret, |
| algorithm="HS256", |
| ) |
| |
| response.set_cookie( |
| self._jwt_cookie_name, |
| value=token, |
| httponly=True, |
| secure=self._jwt_cookie_secure, |
| domain=self._jwt_cookie_domain, |
| samesite=self._jwt_cookie_samesite, |
| ) |
| |
| return response |
| |
| def parse_channel_id_from_request(self, req: Request) -> str: |
| token = req.cookies.get(self._jwt_cookie_name) |
| if not token: |
| raise AsyncQueryTokenException("Token not preset") |
| |
| try: |
| return jwt.decode(token, self._jwt_secret, algorithms=["HS256"])["channel"] |
| except Exception as ex: |
| logger.warning("Parse jwt failed", exc_info=True) |
| raise AsyncQueryTokenException("Failed to parse token") from ex |
| |
| def init_job(self, channel_id: str, user_id: Optional[int]) -> dict[str, Any]: |
| job_id = str(uuid.uuid4()) |
| return build_job_metadata( |
| channel_id, job_id, user_id, status=self.STATUS_PENDING |
| ) |
| |
| # pylint: disable=too-many-arguments |
| def submit_explore_json_job( |
| self, |
| channel_id: str, |
| form_data: dict[str, Any], |
| response_type: str, |
| force: Optional[bool] = False, |
| user_id: Optional[int] = None, |
| ) -> dict[str, Any]: |
| # pylint: disable=import-outside-toplevel |
| from superset import security_manager |
| |
| job_metadata = self.init_job(channel_id, user_id) |
| self._load_explore_json_into_cache_job.delay( |
| {**job_metadata, "guest_token": guest_user.guest_token} |
| if (guest_user := security_manager.get_current_guest_user_if_guest()) |
| else job_metadata, |
| form_data, |
| response_type, |
| force, |
| ) |
| return job_metadata |
| |
| def submit_chart_data_job( |
| self, |
| channel_id: str, |
| form_data: dict[str, Any], |
| user_id: Optional[int] = None, |
| ) -> dict[str, Any]: |
| # pylint: disable=import-outside-toplevel |
| from superset import security_manager |
| |
| # if it's guest user, we want to pass the guest token to the celery task |
| # chart data cache key is calculated based on the current user |
| # this way we can keep the cache key consistent between sync and async command |
| # so that it can be looked up consistently |
| job_metadata = self.init_job(channel_id, user_id) |
| self._load_chart_data_into_cache_job.delay( |
| {**job_metadata, "guest_token": guest_user.guest_token} |
| if (guest_user := security_manager.get_current_guest_user_if_guest()) |
| else job_metadata, |
| form_data, |
| ) |
| return job_metadata |
| |
| def read_events( |
| self, channel: str, last_id: Optional[str] |
| ) -> list[Optional[dict[str, Any]]]: |
| stream_name = f"{self._stream_prefix}{channel}" |
| start_id = increment_id(last_id) if last_id else "-" |
| results = self._redis.xrange(stream_name, start_id, "+", self.MAX_EVENT_COUNT) |
| return [] if not results else list(map(parse_event, results)) |
| |
| def update_job( |
| self, job_metadata: dict[str, Any], status: str, **kwargs: Any |
| ) -> None: |
| if "channel_id" not in job_metadata: |
| raise AsyncQueryJobException("No channel ID specified") |
| |
| if "job_id" not in job_metadata: |
| raise AsyncQueryJobException("No job ID specified") |
| |
| updates = {"status": status, **kwargs} |
| event_data = {"data": json.dumps({**job_metadata, **updates})} |
| |
| full_stream_name = f"{self._stream_prefix}full" |
| scoped_stream_name = f"{self._stream_prefix}{job_metadata['channel_id']}" |
| |
| logger.debug("********** logging event data to stream %s", scoped_stream_name) |
| logger.debug(event_data) |
| |
| self._redis.xadd(scoped_stream_name, event_data, "*", self._stream_limit) |
| self._redis.xadd(full_stream_name, event_data, "*", self._stream_limit_firehose) |
