From 6c8c4b9bef54c16141ae4f4df6cce6ab57075545 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Tue, 18 Aug 2026 09:48:41 -0700 Subject: [PATCH 01/13] chore: introduce coordination service --- UPDATING.md | 1 + docs/admin_docs/configuration/cache.mdx | 7 +- superset/async_events/async_query_manager.py | 87 +-- superset/commands/distributed_lock/acquire.py | 16 +- superset/commands/distributed_lock/base.py | 17 +- superset/commands/distributed_lock/release.py | 17 +- superset/config.py | 41 +- superset/coordination/__init__.py | 437 +++++++++++++++ superset/coordination/exceptions.py | 31 ++ superset/tasks/context.py | 4 +- superset/tasks/manager.py | 517 ++---------------- .../async_events/api_tests.py | 12 +- .../tasks/async_queries_tests.py | 42 +- .../async_events/async_query_manager_tests.py | 107 ++-- tests/unit_tests/coordination/__init__.py | 16 + tests/unit_tests/coordination/test_service.py | 225 ++++++++ .../distributed_lock_tests.py | 87 ++- tests/unit_tests/tasks/test_manager.py | 477 +++++----------- 18 files changed, 1107 insertions(+), 1034 deletions(-) create mode 100644 superset/coordination/__init__.py create mode 100644 superset/coordination/exceptions.py create mode 100644 tests/unit_tests/coordination/__init__.py create mode 100644 tests/unit_tests/coordination/test_service.py diff --git a/UPDATING.md b/UPDATING.md index 11de4acf7f2f..dba1f305c559 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -61,6 +61,7 @@ the old counter to use the outcome-specific replacements. - [42393](https://github.com/apache/superset/pull/42393): Exported dataset YAML now carries a `uuid` for each metric and column so that custom folder assignments (which reference metrics/columns by UUID) survive an import into another workspace. This affects any export bundle that contains datasets, not just a dataset export: chart, dashboard, database and full-asset exports all embed the same dataset YAML, so a dashboard exported from this release also fails to import into an older one even though no dataset was exported directly. As with `folders` and `currency_code_column`, the affected `datasets/` files fail schema validation (`Unknown field: uuid`) when imported into Superset releases that predate this change; regenerate or hand-edit exports for older targets in mixed-version fleets. - [42300](https://github.com/apache/superset/pull/42300): Timeseries charts (line/area/bar) with a Y-axis bound in effect — either an explicit `yAxisBounds` or one derived from `truncateYAxis` — now clamp out-of-range data points to that bound instead of letting ECharts drop the point (and the line segments around it) entirely. Any existing chart with a configured Y-axis bound and data outside it will look different after upgrading: a gap becomes a point pinned to the boundary. The clamp also rewrites the value ECharts reads for that point's tooltip and data label, so the displayed value is the bound rather than the true observation. - [42087](https://github.com/apache/superset/pull/42087): Stored calculated-column and metric expressions are validated when a query is built, under the same sub-query policy already applied to adhoc expressions. Previously only the dataset update path checked them on save, so expressions written by v1 import, by dataset duplication, or before that check existed were never validated. Since `ALLOW_ADHOC_SUBQUERY` defaults to `False` (see [19242](https://github.com/apache/superset/pull/19242)), a dataset whose stored expression contains a sub-query works before upgrading and afterwards fails at chart render with `Custom SQL fields cannot contain sub-queries.` There is no migration step, and the error does not name the offending dataset column, so audit stored expressions before upgrading: either rewrite them without the sub-query, or set `ALLOW_ADHOC_SUBQUERY = True` to keep the previous behaviour for both stored and adhoc expressions. +- The `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` config key is **deprecated** in favor of the shared coordination backend `DISTRIBUTED_COORDINATION_CONFIG`, which powers a single Redis connection for distributed locks, pub/sub, and the async-events streams (the GAQ firehose). All parameters previously accepted by `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` are supported identically under `DISTRIBUTED_COORDINATION_CONFIG` (they use the same `RedisCache`/`RedisSentinelCache` backend). Existing configs continue to work: when `DISTRIBUTED_COORDINATION_CONFIG` is unset, `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` is used as a fallback and logs a one-time deprecation warning. Migrate by renaming the key; the deprecated key will be removed in Superset 8.0. ### Selenium support removed — Playwright is now required for screenshots diff --git a/docs/admin_docs/configuration/cache.mdx b/docs/admin_docs/configuration/cache.mdx index efb2587646e9..022a93becbac 100644 --- a/docs/admin_docs/configuration/cache.mdx +++ b/docs/admin_docs/configuration/cache.mdx @@ -345,9 +345,10 @@ DISTRIBUTED_COORDINATION_CONFIG = { } ``` -By default, connections opened for `DISTRIBUTED_COORDINATION_CONFIG` (as well as -`GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`, which uses the same `RedisCache`/`RedisSentinelCache` -backend) have no socket timeout. This can be overridden with `CACHE_REDIS_SOCKET_TIMEOUT` and +By default, connections opened for `DISTRIBUTED_COORDINATION_CONFIG` (as well as the +deprecated `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`, which uses the same +`RedisCache`/`RedisSentinelCache` backend) have no socket timeout. This can be +overridden with `CACHE_REDIS_SOCKET_TIMEOUT` and `CACHE_REDIS_SOCKET_CONNECT_TIMEOUT`, both in seconds: ```python diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index 2642e2b4b055..fc357c7aee75 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -25,12 +25,12 @@ import jwt from flask import Flask, Request, request, Response, session -from flask_caching.backends.base import BaseCache from superset.async_events.cache_backend import ( RedisCacheBackend, RedisSentinelCacheBackend, ) +from superset.coordination import CoordinationService from superset.utils import json from superset.utils.core import get_user_id @@ -87,6 +87,11 @@ def increment_id(entry_id: str) -> str: def get_cache_backend( config: dict[str, Any], ) -> RedisCacheBackend | RedisSentinelCacheBackend: + """Build a coordination backend from the deprecated GAQ cache config. + + DEPRECATED: retained only so the legacy ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` + setting keeps working as a fallback. Prefer ``DISTRIBUTED_COORDINATION_CONFIG``. + """ cache_config = config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}) cache_type = cache_config.get("CACHE_TYPE") @@ -96,7 +101,6 @@ def get_cache_backend( if cache_type == "RedisSentinelCache": return RedisSentinelCacheBackend.from_config(cache_config) - # TODO: Expand cache backend options. raise UnsupportedCacheBackendError("Unsupported cache backend configuration") @@ -113,7 +117,6 @@ class AsyncQueryManager: def __init__(self) -> None: super().__init__() - self._cache: Optional[BaseCache] = None self._stream_prefix: str = "" self._stream_limit: Optional[int] self._stream_limit_firehose: Optional[int] @@ -137,8 +140,17 @@ def init_app(self, app: Flask) -> None: """ ) - self._cache = get_cache_backend(app.config) - logger.debug("Using GAQ Cache backend as %s", type(self._cache).__name__) + if not ( + app.config.get("DISTRIBUTED_COORDINATION_CONFIG") + or app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( + "CACHE_TYPE" + ) + ): + raise UnsupportedCacheBackendError( + "Global async queries require a coordination backend; configure " + "DISTRIBUTED_COORDINATION_CONFIG (GLOBAL_ASYNC_QUERIES_CACHE_BACKEND " + "is deprecated)." + ) if len(app.config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"]) < 32: raise AsyncQueryTokenException( @@ -298,12 +310,10 @@ def _register_cancellable_job( owner without trusting the client-supplied id. Expires with the JWT so it never outlives the job it guards. """ - if not self._cache: - return - self._cache.set( + CoordinationService.set_value( self._job_registry_key(job_id), json.dumps({"channel_id": channel_id, "user_id": user_id}), - ex=self._jwt_expiration_seconds or None, + ttl=self._jwt_expiration_seconds or None, ) def submit_chart_data_job( @@ -337,33 +347,32 @@ def submit_chart_data_job( def read_events( self, channel: str, last_id: Optional[str] ) -> list[Optional[dict[str, Any]]]: - if not self._cache: + if not CoordinationService.is_backend_defined(): raise CacheBackendNotInitialized("Cache backend not initialized") stream_name = f"{self._stream_prefix}{channel}" start_id = increment_id(last_id) if last_id else "-" - results = self._cache.xrange(stream_name, start_id, "+", self.MAX_EVENT_COUNT) - # Decode bytes to strings, decode_responses is not supported at RedisCache and RedisSentinelCache # noqa: E501 - if isinstance(self._cache, (RedisSentinelCacheBackend, RedisCacheBackend)): - decoded_results = [ - ( - event_id.decode("utf-8"), - { - key.decode("utf-8"): value.decode("utf-8") - for key, value in event_data.items() - }, - ) - for event_id, event_data in results - ] - return ( - [] if not decoded_results else list(map(parse_event, decoded_results)) + results = CoordinationService.stream_range( + stream_name, start_id, "+", self.MAX_EVENT_COUNT + ) + # Decode bytes to strings: the coordination Redis backends do not enable + # decode_responses, so stream_range returns raw bytes. + decoded_results = [ + ( + event_id.decode("utf-8"), + { + key.decode("utf-8"): value.decode("utf-8") + for key, value in event_data.items() + }, ) - return [] if not results else list(map(parse_event, results)) + for event_id, event_data in results + ] + return [] if not decoded_results else list(map(parse_event, decoded_results)) def update_job( self, job_metadata: dict[str, Any], status: str, **kwargs: Any ) -> None: - if not self._cache: + if not CoordinationService.is_backend_defined(): raise CacheBackendNotInitialized("Cache backend not initialized") if "channel_id" not in job_metadata: @@ -388,10 +397,14 @@ def update_job( # SIGUSR1 it is about to receive as a cancellation. if status in (self.STATUS_DONE, self.STATUS_ERROR): if job_id := job_metadata.get("job_id"): - self._cache.delete(self._job_registry_key(job_id)) + CoordinationService.delete_value(self._job_registry_key(job_id)) - self._cache.xadd(scoped_stream_name, event_data, "*", self._stream_limit) - self._cache.xadd(full_stream_name, event_data, "*", self._stream_limit_firehose) + CoordinationService.stream_add( + scoped_stream_name, event_data, "*", self._stream_limit + ) + CoordinationService.stream_add( + full_stream_name, event_data, "*", self._stream_limit_firehose + ) def is_job_cancelled(self, job_id: str) -> bool: """ @@ -401,10 +414,8 @@ def is_job_cancelled(self, job_id: str) -> bool: swallowed and treated as "not cancelled" — a Redis blip must never mask the original error (e.g. a genuine timeout) with a connection error. """ - if not self._cache: - return False try: - raw = self._cache.get(self._job_registry_key(job_id)) + raw = CoordinationService.get_value(self._job_registry_key(job_id)) if raw is None: return False return bool(json.loads(raw).get("cancelled")) @@ -428,11 +439,11 @@ def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int]) -> No :raises AsyncQueryJobException: the job is unknown or already terminal :raises AsyncQueryTokenException: the caller does not own the job """ - if not self._cache: + if not CoordinationService.is_backend_defined(): raise CacheBackendNotInitialized("Cache backend not initialized") key = self._job_registry_key(job_id) - raw = self._cache.get(key) + raw = CoordinationService.get_value(key) if raw is None: raise AsyncQueryJobException("Job not found or already completed") @@ -445,11 +456,11 @@ def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int]) -> No # key still exists (``xx``): if the job finished and cleared its record # between the read above and here, don't recreate a stale record or # revoke a task that is already gone — report it as not found instead. - flagged = self._cache.set( + flagged = CoordinationService.set_value( key, json.dumps({**record, "cancelled": True}), - ex=self._jwt_expiration_seconds or None, - xx=True, + ttl=self._jwt_expiration_seconds or None, + if_present=True, ) if not flagged: raise AsyncQueryJobException("Job not found or already completed") diff --git a/superset/commands/distributed_lock/acquire.py b/superset/commands/distributed_lock/acquire.py index 6cd9a8780994..d12a979d8feb 100644 --- a/superset/commands/distributed_lock/acquire.py +++ b/superset/commands/distributed_lock/acquire.py @@ -27,8 +27,8 @@ from superset.commands.distributed_lock.base import ( BaseDistributedLockCommand, get_default_lock_ttl, - get_redis_client, ) +from superset.coordination import CoordinationService from superset.daos.key_value import KeyValueDAO from superset.exceptions import ( AcquireDistributedLockFailedException, @@ -80,21 +80,21 @@ def __init__( self.ttl_seconds = ttl_seconds or get_default_lock_ttl() def run(self) -> None: - if (redis_client := get_redis_client()) is not None: - self._acquire_redis(redis_client) + if CoordinationService.is_backend_defined(): + self._acquire_redis() else: self._acquire_kv() - def _acquire_redis(self, redis_client: Any) -> None: - """Acquire lock using Redis SET NX EX (atomic).""" + def _acquire_redis(self) -> None: + """Acquire lock using the coordination backend's SET NX EX (atomic).""" try: # SET NX EX: Set if not exists, with expiration # Returns True if lock acquired, None if already exists - acquired = redis_client.set( + acquired = CoordinationService.set_value( self.redis_lock_key, "1", - nx=True, - ex=self.ttl_seconds, + ttl=self.ttl_seconds, + if_absent=True, ) if not acquired: diff --git a/superset/commands/distributed_lock/base.py b/superset/commands/distributed_lock/base.py index 55da69f7ebde..a081df14dbba 100644 --- a/superset/commands/distributed_lock/base.py +++ b/superset/commands/distributed_lock/base.py @@ -19,18 +19,14 @@ import logging import uuid -from typing import Any, TYPE_CHECKING +from typing import Any from flask import current_app from superset.commands.base import BaseCommand from superset.distributed_lock.utils import get_key -from superset.extensions import cache_manager from superset.key_value.types import JsonKeyValueCodec, KeyValueResource -if TYPE_CHECKING: - import redis - logger = logging.getLogger(__name__) @@ -39,17 +35,6 @@ def get_default_lock_ttl() -> int: return int(current_app.config.get("DISTRIBUTED_LOCK_DEFAULT_TTL", 30)) -def get_redis_client() -> "redis.Redis[Any] | None": - """ - Get Redis client from distributed coordination if available. - - Returns None if DISTRIBUTED_COORDINATION_CONFIG is not configured, - allowing fallback to database-backed locking. - """ - backend = cache_manager.distributed_coordination - return backend._cache if backend else None - - class BaseDistributedLockCommand(BaseCommand): """Base command for distributed lock operations.""" diff --git a/superset/commands/distributed_lock/release.py b/superset/commands/distributed_lock/release.py index 14f9deb4df57..65221a731a5c 100644 --- a/superset/commands/distributed_lock/release.py +++ b/superset/commands/distributed_lock/release.py @@ -19,15 +19,12 @@ import logging from functools import partial -from typing import Any import redis from sqlalchemy.exc import SQLAlchemyError -from superset.commands.distributed_lock.base import ( - BaseDistributedLockCommand, - get_redis_client, -) +from superset.commands.distributed_lock.base import BaseDistributedLockCommand +from superset.coordination import CoordinationService from superset.daos.key_value import KeyValueDAO from superset.exceptions import ReleaseDistributedLockFailedException from superset.key_value.exceptions import KeyValueDeleteFailedError @@ -45,15 +42,15 @@ class ReleaseDistributedLock(BaseDistributedLockCommand): """ def run(self) -> None: - if (redis_client := get_redis_client()) is not None: - self._release_redis(redis_client) + if CoordinationService.is_backend_defined(): + self._release_redis() else: self._release_kv() - def _release_redis(self, redis_client: Any) -> None: - """Release lock using Redis DELETE.""" + def _release_redis(self) -> None: + """Release lock using the coordination backend's DELETE.""" try: - redis_client.delete(self.redis_lock_key) + CoordinationService.delete_value(self.redis_lock_key) logger.debug("Released Redis lock: %s", self.redis_lock_key) except redis.RedisError as ex: # Log warning but don't raise - TTL will handle cleanup diff --git a/superset/config.py b/superset/config.py index 0e8ee06b8ebc..46c7b3d0c35a 100644 --- a/superset/config.py +++ b/superset/config.py @@ -2941,6 +2941,14 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument # noq # Global async queries cache backend configuration options: # - Set 'CACHE_TYPE' to 'RedisCache' for RedisCacheBackend. # - Set 'CACHE_TYPE' to 'RedisSentinelCache' for RedisSentinelCacheBackend. +# +# DEPRECATED: prefer DISTRIBUTED_COORDINATION_CONFIG, which powers a single +# coordination service for distributed locks, pub/sub, and the async-events +# streams. When DISTRIBUTED_COORDINATION_CONFIG is set it takes precedence; this +# setting is only used as a fallback (with a deprecation warning) and will be +# removed in Superset 8.0. All parameters here are supported identically under +# DISTRIBUTED_COORDINATION_CONFIG (both use the same +# RedisCache/RedisSentinelCache backend). GLOBAL_ASYNC_QUERIES_CACHE_BACKEND = { "CACHE_TYPE": "RedisCache", "CACHE_REDIS_HOST": "localhost", @@ -3223,24 +3231,43 @@ class ExtraAccessQueryFilters(TypedDict, total=False): # These features require Redis primitives unavailable in generic cache backends: # - Pub/Sub: Real-time message broadcasting between workers # - SET NX EX: Atomic lock acquisition with automatic expiration -# - Streams: Persistent ordered event logs (future) +# - Streams: Persistent ordered event logs (e.g. the Global Async Queries firehose) # # When configured, enables: # - Real-time abort/completion notifications for GTF tasks (vs database polling) # - Redis-based distributed locking (vs KeyValueDAO-backed DistributedLock) -# -# Future: This backend will power a higher-level coordination service exposing -# standardized interfaces for distributed locks, pub/sub, and streams — consolidating -# all advanced Redis primitives under a single connection. Global Async Queries -# (GLOBAL_ASYNC_QUERIES_CACHE_BACKEND) will also be migrated to this configuration. +# - Global Async Queries event streams (the async-events / firehose transport) +# +# This backend powers the higher-level coordination service +# (``superset.coordination.CoordinationService``) exposing standardized interfaces +# for distributed locks, pub/sub, and streams under a single connection. Global +# Async Queries use this when configured; the former +# ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` is deprecated and, if set, used only as a +# fallback (with a deprecation warning). +# +# All parameters previously supported by GLOBAL_ASYNC_QUERIES_CACHE_BACKEND are +# supported here (both go through the same RedisCacheBackend/RedisSentinelCacheBackend +# `from_config`): CACHE_REDIS_HOST, CACHE_REDIS_PORT, CACHE_REDIS_USER, +# CACHE_REDIS_PASSWORD, CACHE_REDIS_DB, CACHE_KEY_PREFIX, CACHE_DEFAULT_TIMEOUT, +# CACHE_REDIS_SSL, CACHE_REDIS_SSL_CERTFILE, CACHE_REDIS_SSL_KEYFILE, +# CACHE_REDIS_SSL_CERT_REQS, CACHE_REDIS_SSL_CA_CERTS, CACHE_REDIS_SOCKET_TIMEOUT, +# CACHE_REDIS_SOCKET_CONNECT_TIMEOUT, and for Sentinel CACHE_REDIS_SENTINELS, +# CACHE_REDIS_SENTINEL_MASTER, CACHE_REDIS_SENTINEL_PASSWORD. # # Example with standard Redis: # DISTRIBUTED_COORDINATION_CONFIG: CacheConfig = { # "CACHE_TYPE": "RedisCache", # "CACHE_REDIS_HOST": "localhost", # "CACHE_REDIS_PORT": 6379, -# "CACHE_REDIS_DB": 0, +# "CACHE_REDIS_USER": "", # "CACHE_REDIS_PASSWORD": "", +# "CACHE_REDIS_DB": 0, +# "CACHE_DEFAULT_TIMEOUT": 300, +# "CACHE_REDIS_SSL": False, # True or False +# "CACHE_REDIS_SSL_CERTFILE": None, +# "CACHE_REDIS_SSL_KEYFILE": None, +# "CACHE_REDIS_SSL_CERT_REQS": "required", +# "CACHE_REDIS_SSL_CA_CERTS": None, # } # # Example with Redis Sentinel: diff --git a/superset/coordination/__init__.py b/superset/coordination/__init__.py new file mode 100644 index 000000000000..db6ccb08a438 --- /dev/null +++ b/superset/coordination/__init__.py @@ -0,0 +1,437 @@ +# 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. +"""Centralized coordination service. + +Provides one connection (``DISTRIBUTED_COORDINATION_CONFIG``) and one interface for +the Valkey/Redis coordination primitives Superset relies on: **pub/sub**, **key/value**, +and event **streams**, plus a higher-level **await/notify** layer (``wait_for_signal`` / +``listen_for_signal``) built on top of them. Distributed **locking** is served by +:class:`~superset.distributed_lock.DistributedLock`, which draws on this service's +backend when one is configured and falls back to a database-backed lock otherwise. + +Historically these were wired up independently: the Global Task Framework used +``DISTRIBUTED_COORDINATION_CONFIG`` (pub/sub and locking) while Global Async Queries +used a separate ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` for its streams, and each caller +hand-rolled its own pub/sub-vs-poll wait loops. Consolidating them here keeps the +architecture modular, gives other components (e.g. the extensions framework) a single +reusable coordination surface, and reduces the number of moving parts. The legacy +``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` is still honored as a fallback (with a +deprecation warning) so existing deployments keep working during the transition. +""" + +from __future__ import annotations + +import logging +import threading +import time +from typing import Any, Callable, TYPE_CHECKING, TypeVar + +from superset.coordination.exceptions import CoordinationBackendUnavailableError + +if TYPE_CHECKING: + from superset.async_events.cache_backend import ( + RedisCacheBackend, + RedisSentinelCacheBackend, + ) + + CoordinationBackend = RedisCacheBackend | RedisSentinelCacheBackend + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + +# Poll cadence for the pub/sub wait loop: how long each ``get_message`` blocks +# before the loop re-checks the predicate, the timeout, and the stop flag. Keeps +# stop latency and missed-message recovery bounded to ~1s. +_PUBSUB_TICK_SECONDS = 1.0 + + +class SignalListener: + """Handle for a background listener started by + :meth:`CoordinationService.listen_for_signal`. + + Wraps the daemon thread, its stop flag, and (in pub/sub mode) the subscription. + :meth:`stop` sets the flag and closes the subscription so a thread blocked in + ``get_message`` wakes immediately, then joins. + """ + + def __init__( + self, + thread: threading.Thread, + stop_event: threading.Event, + pubsub: Any = None, + ) -> None: + self._thread = thread + self._stop_event = stop_event + self._pubsub = pubsub + + def stop(self) -> None: + """Signal the listener to stop and wait briefly for the thread to finish.""" + self._stop_event.set() + # Closing the subscription unblocks a thread parked in get_message so + # teardown is near-immediate rather than waiting a full poll tick. + if self._pubsub is not None: + _close_pubsub(self._pubsub) + if self._thread.is_alive(): + self._thread.join(timeout=2.0) + if self._thread.is_alive(): + # Daemon thread: it will be reaped at process exit. Don't block. + logger.warning( + "Signal listener thread %s did not terminate within 2s.", + self._thread.name, + ) + + +class CoordinationService: + """Single entry point for the Valkey/Redis coordination primitives. + + Two layers of API: + + - **Raw primitives** — ``publish``, ``get`` / ``set`` / ``delete``, + ``stream_add`` / ``stream_range``. These are backend-only and have no fallback: + they raise :class:`CoordinationBackendUnavailableError` when no backend is + configured, rather than silently doing nothing. + - **Higher-level await/notify** — ``wait_for_signal`` (blocking) and + ``listen_for_signal`` (background). These combine a pub/sub channel with a + caller-supplied predicate: + when a backend is defined they wake promptly on a published message, and either + way they fall back to polling the predicate. This keeps the pub/sub-vs-poll + boilerplate in one place; callers just supply a channel and a check. + + All methods are class-level: the service is app-global and resolves its backend + from the shared coordination connection on each call. + + Distributed locking is *not* exposed here: it has its own user-facing interface + (:class:`~superset.distributed_lock.DistributedLock`) that uses this service's + backend when one is defined and falls back to a database-backed lock otherwise. + """ + + _legacy_backend: "CoordinationBackend | None" = None + _legacy_warning_emitted: bool = False + + @classmethod + def get_backend(cls) -> "CoordinationBackend | None": + """Resolve the coordination backend. + + Prefers ``DISTRIBUTED_COORDINATION_CONFIG`` (via the cache manager). Falls + back to the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` when only that + is configured, emitting a one-time deprecation warning. Returns ``None`` when + neither is configured. + """ + from superset.extensions import cache_manager + + if (backend := cache_manager.distributed_coordination) is not None: + return backend + return cls._get_legacy_backend() + + @classmethod + def _get_legacy_backend(cls) -> "CoordinationBackend | None": + if cls._legacy_backend is not None: + return cls._legacy_backend + + from flask import current_app + + if not current_app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( + "CACHE_TYPE" + ): + return None + + if not cls._legacy_warning_emitted: + logger.warning( + "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND is deprecated and will be " + "removed in Superset 8.0; configure DISTRIBUTED_COORDINATION_CONFIG " + "instead so a single connection powers distributed locks, pub/sub, " + "and streams." + ) + cls._legacy_warning_emitted = True + + from superset.async_events.async_query_manager import get_cache_backend + + cls._legacy_backend = get_cache_backend(current_app.config) + return cls._legacy_backend + + @classmethod + def is_backend_defined(cls) -> bool: + """Whether a coordination backend is defined. + + Some operations require the Valkey/Redis backend + (``DISTRIBUTED_COORDINATION_CONFIG``) to be configured; those that do note it + on their own docstring. Best-effort callers should branch on this before + invoking a backend-dependent operation instead of catching + :class:`CoordinationBackendUnavailableError`. + """ + return cls.get_backend() is not None + + @classmethod + def _require_backend(cls) -> "CoordinationBackend": + """Return the backend or raise if none is configured. + + Used by the backend-only primitives (pub/sub publish, key/value, streams) + so a missing backend fails loudly instead of silently no-op'ing. + """ + backend = cls.get_backend() + if backend is None: + raise CoordinationBackendUnavailableError( + "No coordination backend configured; set " + "DISTRIBUTED_COORDINATION_CONFIG to enable key/value and stream " + "operations." + ) + return backend + + # -- Pub/Sub ------------------------------------------------------------- + + @classmethod + def publish(cls, channel: str, message: str) -> int: + """Publish a message to a channel; returns the subscriber count. + + Only publishing is offered here — subscribing needs the native connection + (a long-lived subscription with its own receive loop), so consumers that + subscribe should obtain it via :meth:`get_backend`. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().publish(channel, message) + + # -- Key/Value ----------------------------------------------------------- + + @classmethod + def get_value(cls, key: str) -> Any: + """Return the raw (bytes) value at ``key``, or ``None`` if absent. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().get(key) + + @classmethod + def set_value( + cls, + key: str, + value: Any, + ttl: int | None = None, + if_absent: bool = False, + if_present: bool = False, + ) -> bool | None: + """Store ``value`` at ``key``. + + :param ttl: optional expiry, in seconds. + :param if_absent: only set if the key does not already exist. + :param if_present: only set if the key already exists. + :returns: ``True`` on success, or ``None`` when an ``if_absent`` / + ``if_present`` condition prevented the write. + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().set( + key, value, ex=ttl, nx=if_absent, xx=if_present + ) + + @classmethod + def delete_value(cls, *keys: str) -> int: + """Delete one or more keys; returns the number deleted. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().delete(*keys) + + # -- Streams ------------------------------------------------------------- + + @classmethod + def stream_add( + cls, + stream: str, + data: dict[str, Any], + event_id: str = "*", + max_len: int | None = None, + ) -> str: + """Append an event to a stream; returns the generated event id. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().xadd(stream, data, event_id, max_len) + + @classmethod + def stream_range( + cls, + stream: str, + start: str = "-", + end: str = "+", + count: int | None = None, + ) -> list[Any]: + """Read a range of events from a stream. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().xrange(stream, start, end, count) + + # -- Await / notify ------------------------------------------------------ + + @classmethod + def wait_for_signal( + cls, + channel: str, + check: Callable[[], T | None], + *, + timeout: float | None = None, + poll_interval: float = 1.0, + ) -> T: + """Block until ``check()`` returns a non-``None`` value; return that value. + + ``check`` is the source of truth (typically a metastore read). When a + coordination backend is defined, this subscribes to ``channel`` and re-runs + ``check`` promptly whenever a message is published; otherwise it polls + ``check`` every ``poll_interval`` seconds. ``check`` is also re-evaluated on + every tick even in pub/sub mode, so a signal published before the subscription + (or a dropped message) is still caught. + + :param channel: pub/sub channel that peers publish to when the awaited state + is reached (used only as a low-latency wake-up; correctness relies on + ``check``). + :param check: returns a truthy result once the wait is satisfied, else + ``None``. + :param timeout: max seconds to wait; ``None`` waits indefinitely. + :param poll_interval: poll cadence when no backend is defined. + :raises TimeoutError: if ``timeout`` elapses before ``check`` is satisfied. + """ + deadline = None if timeout is None else time.monotonic() + timeout + backend = cls.get_backend() + pubsub = backend.pubsub() if backend is not None else None + try: + if pubsub is not None: + pubsub.subscribe(channel) + while True: + # ``check`` is the source of truth; run it first so the fast path and + # any signal missed before subscribing are both covered. + if (result := check()) is not None: + return result + remaining = ( + None if deadline is None else max(0.0, deadline - time.monotonic()) + ) + if remaining is not None and remaining <= 0: + raise TimeoutError(f"Timed out waiting on channel {channel}") + cls._wait_tick(pubsub, poll_interval, remaining) + finally: + if pubsub is not None: + _close_pubsub(pubsub) + + @staticmethod + def _wait_tick(pubsub: Any, poll_interval: float, remaining: float | None) -> None: + """Block for one wait tick: a pub/sub message (nudge) or a poll sleep.""" + if pubsub is not None: + wait = ( + _PUBSUB_TICK_SECONDS + if remaining is None + else min(_PUBSUB_TICK_SECONDS, remaining) + ) + pubsub.get_message(ignore_subscribe_messages=True, timeout=wait) + else: + time.sleep( + poll_interval if remaining is None else min(poll_interval, remaining) + ) + + @classmethod + def listen_for_signal( + cls, + channel: str, + check: Callable[[], bool], + on_signal: Callable[[], None], + *, + poll_interval: float, + name: str | None = None, + ) -> SignalListener: + """Run a background daemon that invokes ``on_signal`` once ``check`` is true. + + Same wake-vs-poll model as :meth:`wait_for_signal`: a published message on + ``channel`` wakes the loop when a backend is defined, otherwise it polls + ``check`` every ``poll_interval`` seconds. The thread stops after firing + ``on_signal`` once, or when :meth:`SignalListener.stop` is called. + + :param channel: pub/sub channel peers publish to when the condition is met. + :param check: returns ``True`` once ``on_signal`` should fire. + :param on_signal: invoked (once) when ``check`` becomes true. + :param poll_interval: poll cadence when no backend is defined. + :param name: optional thread name suffix for logging. + """ + stop_event = threading.Event() + backend = cls.get_backend() + pubsub = backend.pubsub() if backend is not None else None + if pubsub is not None: + # Subscribe in the caller's thread so a connection failure surfaces here + # (fail-fast) rather than dying silently in the daemon thread. + try: + pubsub.subscribe(channel) + except Exception: + _close_pubsub(pubsub) + raise + thread = threading.Thread( + target=cls._run_listen_loop, + args=(channel, check, on_signal, stop_event, poll_interval, pubsub), + daemon=True, + name=f"coord-listen-{name or channel}", + ) + thread.start() + return SignalListener(thread, stop_event, pubsub) + + @classmethod + def _run_listen_loop( + cls, + channel: str, + check: Callable[[], bool], + on_signal: Callable[[], None], + stop_event: threading.Event, + poll_interval: float, + pubsub: Any, + ) -> None: + """Body of the background listener thread (see :meth:`listen`).""" + try: + while not stop_event.is_set(): + try: + if check(): + on_signal() + return + if pubsub is not None: + # Blocks up to a tick; the message is just a wake-up nudge. + pubsub.get_message( + ignore_subscribe_messages=True, + timeout=_PUBSUB_TICK_SECONDS, + ) + else: + stop_event.wait(timeout=poll_interval) + except (ValueError, OSError) as ex: + # Connection torn down (e.g. stop() closing the subscription, or + # shutdown). Expected when stopping; otherwise surface it and bail. + if not stop_event.is_set(): + logger.error( + "Signal listener on %s failed: %s", + channel, + ex, + exc_info=True, + ) + return + except Exception: # pylint: disable=broad-except + if not stop_event.is_set(): + logger.exception("Signal listener on %s crashed", channel) + finally: + if pubsub is not None: + _close_pubsub(pubsub) + + +def _close_pubsub(pubsub: Any) -> None: + """Best-effort unsubscribe + close of a pub/sub subscription.""" + try: + pubsub.unsubscribe() + pubsub.close() + except Exception as ex: # pylint: disable=broad-except + logger.debug("Error closing pub/sub subscription: %s", ex) diff --git a/superset/coordination/exceptions.py b/superset/coordination/exceptions.py new file mode 100644 index 000000000000..00e2b794afdf --- /dev/null +++ b/superset/coordination/exceptions.py @@ -0,0 +1,31 @@ +# 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. +"""Exceptions raised by the coordination service.""" + +from __future__ import annotations + + +class CoordinationBackendUnavailableError(Exception): + """Raised when a Valkey/Redis-only primitive is used without a backend. + + Pub/sub, streams, and key/value operations have no in-service fallback, so + calling them without a configured coordination backend is a programming or + configuration error rather than a silently-ignored no-op. Callers that have + their own fallback (e.g. database polling) should gate on + :meth:`superset.coordination.CoordinationService.is_backend_defined` instead of + catching this. + """ diff --git a/superset/tasks/context.py b/superset/tasks/context.py index 76c8a01b8550..b4052022e558 100644 --- a/superset/tasks/context.py +++ b/superset/tasks/context.py @@ -35,8 +35,8 @@ from superset.tasks.utils import progress_update if TYPE_CHECKING: + from superset.coordination import SignalListener from superset.models.tasks import Task - from superset.tasks.manager import AbortListener logger = logging.getLogger(__name__) @@ -67,7 +67,7 @@ def __init__(self, task: "Task") -> None: self._task_uuid = task.uuid self._cleanup_handlers: list[Callable[[], None]] = [] self._abort_handlers: list[Callable[[], None]] = [] - self._abort_listener: "AbortListener | None" = None + self._abort_listener: "SignalListener | None" = None self._abort_detected = False self._abort_handlers_completed = False # Track if all abort handlers finished self._execution_completed = False # Set by executor after task work completes diff --git a/superset/tasks/manager.py b/superset/tasks/manager.py index 08adb639a9ce..b1ab21c98da7 100644 --- a/superset/tasks/manager.py +++ b/superset/tasks/manager.py @@ -19,8 +19,6 @@ from __future__ import annotations import logging -import threading -import time from typing import Any, Callable, TYPE_CHECKING from uuid import UUID @@ -28,72 +26,18 @@ from flask import has_app_context from superset_core.tasks.types import TaskProperties, TaskScope -from superset.async_events.cache_backend import ( - RedisCacheBackend, - RedisSentinelCacheBackend, -) -from superset.extensions import cache_manager from superset.tasks.constants import ABORT_STATES, TERMINAL_STATES from superset.tasks.utils import generate_random_task_key if TYPE_CHECKING: from flask import Flask + from superset.coordination import SignalListener from superset.models.tasks import Task logger = logging.getLogger(__name__) -class AbortListener: - """ - Handle for a background abort listener. - - Returned by TaskManager.listen_for_abort() to allow stopping the listener. - """ - - def __init__( - self, - task_uuid: UUID, - thread: threading.Thread, - stop_event: threading.Event, - pubsub: redis.client.PubSub | None = None, - ) -> None: - self._task_uuid = task_uuid - self._thread = thread - self._stop_event = stop_event - self._pubsub = pubsub - - def stop(self) -> None: - """Stop the abort listener.""" - self._stop_event.set() - - # Close pub/sub subscription if active - if self._pubsub is not None: - try: - self._pubsub.unsubscribe() - self._pubsub.close() - except Exception as ex: - logger.debug("Error closing pub/sub during stop: %s", ex) - - # Wait for thread to finish (with timeout to avoid blocking indefinitely) - if self._thread.is_alive(): - self._thread.join(timeout=2.0) - - # Check if thread is still running after timeout - if self._thread.is_alive(): - # Thread is a daemon, so it will be killed when process exits. - # Log warning but continue - cleanup will still proceed. - logger.warning( - "Abort listener thread for task %s did not terminate within " - "2 seconds. Thread will be terminated when process exits.", - self._task_uuid, - ) - else: - logger.debug("Stopped abort listener for task %s", self._task_uuid) - else: - logger.debug("Stopped abort listener for task %s", self._task_uuid) - - class TaskManager: """ Handles task creation, scheduling, and abort notifications. @@ -132,24 +76,6 @@ def init_app(cls, app: Flask) -> None: cls._initialized = True - @classmethod - def _get_cache(cls) -> RedisCacheBackend | RedisSentinelCacheBackend | None: - """ - Get the distributed coordination backend. - - :returns: The distributed coordination backend, or None if not configured - """ - return cache_manager.distributed_coordination - - @classmethod - def is_pubsub_available(cls) -> bool: - """ - Check if Redis pub/sub backend is configured and available. - - :returns: True if Redis is available for pub/sub, False otherwise - """ - return cls._get_cache() is not None - @classmethod def get_abort_channel(cls, task_uuid: UUID) -> str: """ @@ -168,13 +94,14 @@ def publish_abort(cls, task_uuid: UUID) -> bool: :param task_uuid: UUID of the task to abort :returns: True if message was published, False if Redis unavailable """ - cache = cls._get_cache() - if not cache: + from superset.coordination import CoordinationService + + if not CoordinationService.is_backend_defined(): return False try: channel = cls.get_abort_channel(task_uuid) - subscriber_count = cache.publish(channel, "abort") + subscriber_count = CoordinationService.publish(channel, "abort") logger.debug( "Published abort to channel %s (%d subscribers)", channel, @@ -207,13 +134,14 @@ def publish_completion(cls, task_uuid: UUID, status: str) -> bool: :param status: Final status of the task :returns: True if message was published, False if Redis unavailable """ - cache = cls._get_cache() - if not cache: + from superset.coordination import CoordinationService + + if not CoordinationService.is_backend_defined(): return False try: channel = cls.get_completion_channel(task_uuid) - subscriber_count = cache.publish(channel, status) + subscriber_count = CoordinationService.publish(channel, status) logger.debug( "Published completion to channel %s (status=%s, %d subscribers)", channel, @@ -236,8 +164,10 @@ def wait_for_completion( """ Block until task reaches terminal state. - Uses Redis pub/sub if configured for low-latency, low-CPU waiting. - Uses database polling if Redis is not configured. + Delegates the pub/sub-wake-else-poll orchestration to + :meth:`CoordinationService.wait_for_signal`; here we only supply the + completion channel and a metastore predicate that returns the task once it is + terminal. :param task_uuid: UUID of the task to wait for :param timeout: Maximum time to wait in seconds (None = no limit) @@ -247,17 +177,9 @@ def wait_for_completion( :raises TimeoutError: If timeout expires before task completes :raises ValueError: If task not found """ + from superset.coordination import CoordinationService from superset.daos.tasks import TaskDAO - start_time = time.monotonic() - - def time_remaining() -> float | None: - if timeout is None: - return None - elapsed = time.monotonic() - start_time - remaining = timeout - elapsed - return remaining if remaining > 0 else 0 - def get_task() -> "Task | None": # Reads back the task named by the caller's own task_uuid, not # a user-requested lookup; see TaskFilter for the @@ -269,135 +191,20 @@ def get_task() -> "Task | None": ) return TaskDAO.find_one_or_none(uuid=task_uuid, skip_base_filter=True) - # Check current state first - task = get_task() - if not task: + # Fail fast if the task doesn't exist at all. + if get_task() is None: raise ValueError(f"Task {task_uuid} not found") - if task.status in TERMINAL_STATES: - return task - - logger.debug( - "Waiting for task %s to complete (current status=%s, timeout=%s)", - task_uuid, - task.status, - timeout, - ) - - # Use Redis pub/sub if configured - if (cache := cls._get_cache()) is not None: - task = cls._wait_via_pubsub( - task_uuid, - cache.pubsub(), - timeout, - poll_interval, - get_task, - time_remaining, - ) - if task: - return task - # Should not reach here - _wait_via_pubsub returns task or raises - raise RuntimeError(f"Unexpected state waiting for task {task_uuid}") - - # Use database polling when Redis is not configured - return cls._wait_via_polling(task_uuid, poll_interval, get_task, time_remaining) - - @classmethod - def _wait_via_pubsub( - cls, - task_uuid: UUID, - pubsub: redis.client.PubSub, - timeout: float | None, - poll_interval: float, - get_task: Callable[[], "Task | None"], - time_remaining: Callable[[], float | None], - ) -> "Task | None": - """ - Wait for task completion using Redis pub/sub. - - :returns: Task when completed - :raises TimeoutError: If timeout expires - :raises redis.RedisError: If Redis connection fails - """ - channel = cls.get_completion_channel(task_uuid) - pubsub.subscribe(channel) - - try: - while True: - remaining = time_remaining() - if remaining is not None and remaining <= 0: - raise TimeoutError( - f"Timeout waiting for task {task_uuid} to complete" - ) - - # Wait for message with short timeout for responsive checking - wait_time = min(1.0, remaining) if remaining else 1.0 - message = pubsub.get_message( - ignore_subscribe_messages=True, - timeout=wait_time, - ) - - if message and message.get("type") == "message": - # Completion received - fetch fresh task state - logger.debug( - "Received completion message for task %s: %s", - task_uuid, - message.get("data"), - ) - task = get_task() - if task and task.status in TERMINAL_STATES: - return task - - # Also check database periodically in case we missed the message - # (e.g., task completed before we subscribed) - task = get_task() - if task and task.status in TERMINAL_STATES: - logger.debug( - "Task %s completed (detected via db check): status=%s", - task_uuid, - task.status, - ) - return task - - finally: - pubsub.unsubscribe() - pubsub.close() - - @classmethod - def _wait_via_polling( - cls, - task_uuid: UUID, - poll_interval: float, - get_task: Callable[[], "Task | None"], - time_remaining: Callable[[], float | None], - ) -> "Task": - """ - Wait for task completion using database polling. - - :returns: Task when completed - :raises TimeoutError: If timeout expires - :raises ValueError: If task not found - """ - while True: - remaining = time_remaining() - if remaining is not None and remaining <= 0: - raise TimeoutError(f"Timeout waiting for task {task_uuid} to complete") - + def terminal_task() -> "Task | None": task = get_task() - if not task: - raise ValueError(f"Task {task_uuid} not found") + return task if task and task.status in TERMINAL_STATES else None - if task.status in TERMINAL_STATES: - logger.debug( - "Task %s completed (detected via polling): status=%s", - task_uuid, - task.status, - ) - return task - - # Sleep with timeout awareness - sleep_time = min(poll_interval, remaining) if remaining else poll_interval - time.sleep(sleep_time) + return CoordinationService.wait_for_signal( + cls.get_completion_channel(task_uuid), + terminal_task, + timeout=timeout, + poll_interval=poll_interval, + ) @classmethod def listen_for_abort( @@ -406,72 +213,39 @@ def listen_for_abort( callback: Callable[[], None], poll_interval: float, app: Any = None, - ) -> AbortListener: + ) -> "SignalListener": """ Start listening for abort notifications for a task. - Uses Redis pub/sub if configured, otherwise uses database polling. - The callback is invoked when an abort is detected. + Delegates the pub/sub-wake-else-poll orchestration to + :meth:`CoordinationService.listen_for_signal`; here we only supply the abort + channel and an (app-context-aware) abort predicate + callback. :param task_uuid: UUID of the task to monitor (native UUID) :param callback: Function to call when abort is detected :param poll_interval: Interval for database polling (when Redis not configured) :param app: Flask app for database access in background thread - :returns: AbortListener handle to stop listening - """ - stop_event = threading.Event() - pubsub: redis.client.PubSub | None = None - uuid_str = str(task_uuid) - - # Use Redis pub/sub if configured - if (cache := cls._get_cache()) is not None: - pubsub = cache.pubsub() - channel = cls.get_abort_channel(task_uuid) - pubsub.subscribe(channel) - logger.debug("Subscribed to abort channel: %s", channel) - - # Start pub/sub listener thread - thread = threading.Thread( - target=cls._listen_pubsub, - args=(task_uuid, pubsub, callback, stop_event, app), - daemon=True, - name=f"abort-listener-{uuid_str[:8]}", - ) - logger.debug("Started pub/sub abort listener for task %s", task_uuid) - else: - # Use polling when Redis is not configured - pubsub = None - thread = threading.Thread( - target=cls._poll_for_abort, - args=(task_uuid, callback, stop_event, poll_interval, app), - daemon=True, - name=f"abort-poller-{uuid_str[:8]}", - ) - logger.debug( - "Started database abort polling for task %s (interval=%ss)", - task_uuid, - poll_interval, - ) - - thread.start() - return AbortListener(task_uuid, thread, stop_event, pubsub) - - @staticmethod - def _invoke_callback_with_context( - callback: Callable[[], None], - app: Any, - ) -> None: - """ - Invoke callback with Flask app context if provided. - - :param callback: Function to invoke - :param app: Flask app for context, or None + :returns: SignalListener handle to stop listening """ - if app and not has_app_context(): - with app.app_context(): - callback() - else: - callback() + from superset.coordination import CoordinationService + + def in_context(fn: Callable[[], Any]) -> Callable[[], Any]: + # The listener runs in a background thread; DB access needs app context. + def wrapped() -> Any: + if app and not has_app_context(): + with app.app_context(): + return fn() + return fn() + + return wrapped + + return CoordinationService.listen_for_signal( + cls.get_abort_channel(task_uuid), + check=in_context(lambda: cls._check_abort_status(task_uuid)), + on_signal=in_context(callback), + poll_interval=poll_interval, + name=str(task_uuid), + ) @classmethod def _check_abort_status(cls, task_uuid: UUID) -> bool: @@ -488,201 +262,6 @@ def _check_abort_status(cls, task_uuid: UUID) -> bool: task = TaskDAO.find_one_or_none(uuid=task_uuid, skip_base_filter=True) return task is not None and task.status in ABORT_STATES - @classmethod - def _run_abort_listener_loop( - cls, - task_uuid: UUID, - callback: Callable[[], None], - stop_event: threading.Event, - interval: float, - app: Any, - check_fn: Callable[[], bool], - source: str, - ) -> None: - """ - Common abort listener loop used by both pub/sub and polling modes. - - :param task_uuid: UUID of the task to monitor (native UUID) - :param callback: Function to call when abort is detected - :param stop_event: Event to signal loop termination - :param interval: Wait interval between checks - :param app: Flask app for context - :param check_fn: Function that returns True if abort was detected - :param source: Source identifier for logging ("pub/sub" or "polling") - """ - while not stop_event.is_set(): - try: - if check_fn(): - logger.info( - "Abort detected via %s for task %s", - source, - task_uuid, - ) - cls._invoke_callback_with_context(callback, app) - break - - # Wait for interval or until stop is requested - stop_event.wait(timeout=interval) - - except (ValueError, OSError) as ex: - # ValueError/OSError with "I/O operation on closed file" or - # "Bad file descriptor" typically means the connection was closed - # during shutdown. Check if stop was requested. - if stop_event.is_set(): - logger.debug( - "Abort %s for task %s stopped cleanly (connection closed)", - source, - task_uuid, - ) - else: - logger.error( - "Error in abort %s for task %s: %s", - source, - task_uuid, - str(ex), - exc_info=True, - ) - break - - except Exception as ex: - # Check if stop was requested - if so, this may be expected - if stop_event.is_set(): - logger.debug( - "Abort %s for task %s stopped with exception: %s", - source, - task_uuid, - ex, - ) - else: - logger.error( - "Error in abort %s for task %s: %s", - source, - task_uuid, - str(ex), - exc_info=True, - ) - break - - @classmethod - def _listen_pubsub( - cls, - task_uuid: UUID, - pubsub: redis.client.PubSub, - callback: Callable[[], None], - stop_event: threading.Event, - app: Any, - ) -> None: - """Listen for abort via Redis pub/sub.""" - # Track if abort was received to avoid double-callback - abort_received = False - - def check_pubsub() -> bool: - nonlocal abort_received - message = pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0) - if message is not None and message.get("type") == "message": - abort_received = True - return True - return False - - try: - cls._run_abort_listener_loop( - task_uuid=task_uuid, - callback=callback, - stop_event=stop_event, - interval=0, # pub/sub has its own timeout in get_message - app=app, - check_fn=check_pubsub, - source="pub/sub", - ) - - except redis.RedisError as ex: - # Check if we were asked to stop - if so, this is expected - if stop_event.is_set(): - logger.debug( - "Abort listener for task %s stopped (Redis error: %s)", - task_uuid, - ex, - ) - else: - # Log error but don't fall back - let the failure be visible - logger.error( - "Redis signal backend failed for task %s abort listener: %s. " - "Task may not receive abort signal.", - task_uuid, - ex, - ) - - except (ValueError, OSError) as ex: - # ValueError: "I/O operation on closed file" - expected when stop() closes - # OSError: Similar connection-closed errors - if stop_event.is_set(): - # Clean shutdown, expected behavior - logger.debug( - "Abort listener for task %s stopped cleanly", - task_uuid, - ) - else: - # Unexpected error while running - logger.error( - "Error in abort listener for task %s: %s", - task_uuid, - str(ex), - exc_info=True, - ) - - except Exception as ex: - # Only log as error if we weren't asked to stop - if stop_event.is_set(): - logger.debug( - "Abort listener for task %s stopped with exception: %s", - task_uuid, - ex, - ) - else: - logger.error( - "Error in abort listener for task %s: %s", - task_uuid, - str(ex), - exc_info=True, - ) - - finally: - # Clean up pub/sub subscription - try: - pubsub.unsubscribe() - pubsub.close() - except Exception as ex: - logger.debug("Error closing pub/sub during cleanup: %s", ex) - - @classmethod - def _poll_for_abort( - cls, - task_uuid: UUID, - callback: Callable[[], None], - stop_event: threading.Event, - interval: float, - app: Any, - ) -> None: - """Background polling loop - used when Redis pub/sub is not configured.""" - - def check_database() -> bool: - # Need app context for database access - if app and not has_app_context(): - with app.app_context(): - return cls._check_abort_status(task_uuid) - else: - return cls._check_abort_status(task_uuid) - - cls._run_abort_listener_loop( - task_uuid=task_uuid, - callback=callback, - stop_event=stop_event, - interval=interval, - app=app, - check_fn=check_database, - source="polling", - ) - @staticmethod def submit_task( task_type: str, diff --git a/tests/integration_tests/async_events/api_tests.py b/tests/integration_tests/async_events/api_tests.py index 5be8022d5b7f..bceca0845fbe 100644 --- a/tests/integration_tests/async_events/api_tests.py +++ b/tests/integration_tests/async_events/api_tests.py @@ -48,14 +48,16 @@ def run_test_with_cache_backend(self, cache_backend_cls: Type[Any], test_func): app._got_first_request = False async_query_manager_factory.init_app(app) - # Create a mock cache backend instance + # The manager resolves its backend through CoordinationService.get_backend(), + # so patch there rather than assigning a (now-removed) private attribute. mock_cache = mock.Mock(spec=cache_backend_cls) - # Set the mock cache instance - async_query_manager._cache = mock_cache - self.login(ADMIN_USERNAME) - test_func(mock_cache) + with mock.patch( + "superset.coordination.CoordinationService.get_backend", + return_value=mock_cache, + ): + test_func(mock_cache) def _test_events_logic(self, mock_cache): with mock.patch.object(mock_cache, "xrange") as mock_xrange: diff --git a/tests/integration_tests/tasks/async_queries_tests.py b/tests/integration_tests/tasks/async_queries_tests.py index 16d618f0ece2..f254840804d4 100644 --- a/tests/integration_tests/tasks/async_queries_tests.py +++ b/tests/integration_tests/tasks/async_queries_tests.py @@ -21,12 +21,7 @@ import pytest from celery.exceptions import SoftTimeLimitExceeded -from parameterized import parameterized -from superset.async_events.cache_backend import ( - RedisCacheBackend, - RedisSentinelCacheBackend, -) from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.commands.chart.exceptions import ChartDataQueryFailedError from superset.extensions import async_query_manager, security_manager @@ -46,24 +41,13 @@ "load_birth_names_data", "load_birth_names_dashboard_with_slices" ) class TestAsyncQueries(SupersetTestCase): - @parameterized.expand( - [ - ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)), - ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)), - ] - ) @mock.patch("superset.tasks.async_queries.set_form_data") @mock.patch.object(async_query_manager, "update_job") - def test_load_chart_data_into_cache( - self, cache_type, cache_backend, mock_update_job, mock_set_form_data - ): + def test_load_chart_data_into_cache(self, mock_update_job, mock_set_form_data): from superset.tasks.async_queries import load_chart_data_into_cache app._got_first_request = False - async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend) - async_query_manager.init_app(app) - query_context = get_query_context("birth_names") user = security_manager.find_user("gamma") job_metadata = { @@ -81,26 +65,15 @@ def test_load_chart_data_into_cache( job_metadata, "done", result_url=mock.ANY ) - @parameterized.expand( - [ - ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)), - ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)), - ] - ) @mock.patch.object( ChartDataCommand, "run", side_effect=ChartDataQueryFailedError("Error: foo") ) @mock.patch.object(async_query_manager, "update_job") - def test_load_chart_data_into_cache_error( - self, cache_type, cache_backend, mock_update_job, mock_run_command - ): + def test_load_chart_data_into_cache_error(self, mock_update_job, mock_run_command): from superset.tasks.async_queries import load_chart_data_into_cache app._got_first_request = False - async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend) - async_query_manager.init_app(app) - query_context = get_query_context("birth_names") user = security_manager.find_user("gamma") job_metadata = { @@ -117,24 +90,15 @@ def test_load_chart_data_into_cache_error( errors = [{"message": "Error: foo"}] mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors) - @parameterized.expand( - [ - ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)), - ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)), - ] - ) @mock.patch.object(ChartDataCommand, "run") @mock.patch.object(async_query_manager, "update_job") def test_soft_timeout_load_chart_data_into_cache( - self, cache_type, cache_backend, mock_update_job, mock_run_command + self, mock_update_job, mock_run_command ): from superset.tasks.async_queries import load_chart_data_into_cache app._got_first_request = False - async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend) - async_query_manager.init_app(app) - user = security_manager.find_user("gamma") form_data = {} job_metadata = { diff --git a/tests/unit_tests/async_events/async_query_manager_tests.py b/tests/unit_tests/async_events/async_query_manager_tests.py index edd896124749..f6aaa972ba16 100644 --- a/tests/unit_tests/async_events/async_query_manager_tests.py +++ b/tests/unit_tests/async_events/async_query_manager_tests.py @@ -20,7 +20,7 @@ from flask import g from jwt import encode -from pytest import fixture, mark, raises # noqa: PT013 +from pytest import fixture, raises # noqa: PT013 from superset import security_manager from superset.async_events.async_query_manager import ( @@ -30,7 +30,6 @@ ) from superset.async_events.cache_backend import ( RedisCacheBackend, - RedisSentinelCacheBackend, ) from superset.utils import json @@ -263,23 +262,13 @@ def test_parse_channel_id_from_request_as_guest_user_differs_per_scope( assert with_datasets != with_rev -@mark.parametrize( - "cache_type, cache_backend", - [ - ("RedisCacheBackend", mock.Mock(spec=RedisCacheBackend)), - ("RedisSentinelCacheBackend", mock.Mock(spec=RedisSentinelCacheBackend)), - ], -) @mock.patch("superset.is_feature_enabled") def test_submit_chart_data_job_as_guest_user( - is_feature_enabled_mock, async_query_manager, cache_type, cache_backend + is_feature_enabled_mock, async_query_manager ): is_feature_enabled_mock.return_value = True set_current_as_guest_user() - # Mock the get_cache_backend method to return the current cache backend - async_query_manager.get_cache_backend = mock.Mock(return_value=cache_backend) - job_mock = Mock() async_query_manager._load_chart_data_into_cache_job = job_mock job_meta = async_query_manager.submit_chart_data_job( @@ -370,32 +359,44 @@ def test_view(): @fixture -def cancellable_manager(): +def coordination_backend(): + """Patch the coordination service to a mock Redis backend and expose it.""" + backend = mock.Mock(spec=RedisCacheBackend) + with mock.patch( + "superset.coordination.CoordinationService.get_backend", + return_value=backend, + ): + yield backend + + +@fixture +def cancellable_manager(coordination_backend): """A manager wired to a mock Redis backend for cancellation tests.""" manager = AsyncQueryManager() manager._jwt_expiration_seconds = 3600 manager._stream_prefix = "async-events-" - manager._cache = mock.Mock(spec=RedisCacheBackend) return manager -def test_init_job_registers_cancellable_record(cancellable_manager): +def test_init_job_registers_cancellable_record( + cancellable_manager, coordination_backend +): """init_job persists the owner identity a later cancel must match.""" cancellable_manager.init_job("chan-1", 7) - cancellable_manager._cache.set.assert_called_once() - key, value = cancellable_manager._cache.set.call_args.args + coordination_backend.set.assert_called_once() + key, value = coordination_backend.set.call_args.args assert key.startswith("async-events-job-cancel:") assert json.loads(value) == {"channel_id": "chan-1", "user_id": 7} -def test_cancel_job_authorized_revokes_task(cancellable_manager): +def test_cancel_job_authorized_revokes_task(cancellable_manager, coordination_backend): cancellable_manager._stream_limit = 100 cancellable_manager._stream_limit_firehose = 1000 - cancellable_manager._cache.get.return_value = json.dumps( + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7} ) - cancellable_manager._cache.set.return_value = True + coordination_backend.set.return_value = True with mock.patch("superset.extensions.celery_app") as celery_app: cancellable_manager.cancel_job("job-1", "chan-1", 7) @@ -405,26 +406,24 @@ def test_cancel_job_authorized_revokes_task(cancellable_manager): ) # The job is flagged cancelled (conditionally, xx=True) so the worker knows # what the signal it is about to receive means. - assert cancellable_manager._cache.set.call_args.kwargs["xx"] is True - flagged = json.loads(cancellable_manager._cache.set.call_args.args[1]) + assert coordination_backend.set.call_args.kwargs["xx"] is True + flagged = json.loads(coordination_backend.set.call_args.args[1]) assert flagged["cancelled"] is True -def test_cancel_job_emits_the_terminal_event(cancellable_manager): +def test_cancel_job_emits_the_terminal_event(cancellable_manager, coordination_backend): """A task revoked before a worker picks it up never reports on itself.""" cancellable_manager._stream_limit = 100 cancellable_manager._stream_limit_firehose = 1000 - cancellable_manager._cache.get.return_value = json.dumps( + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7} ) - cancellable_manager._cache.set.return_value = True + coordination_backend.set.return_value = True with mock.patch("superset.extensions.celery_app"): cancellable_manager.cancel_job("job-1", "chan-1", 7) - scoped_stream, event_data = cancellable_manager._cache.xadd.call_args_list[0].args[ - :2 - ] + scoped_stream, event_data = coordination_backend.xadd.call_args_list[0].args[:2] assert scoped_stream == "async-events-chan-1" assert json.loads(event_data["data"]) == { "channel_id": "chan-1", @@ -436,16 +435,18 @@ def test_cancel_job_emits_the_terminal_event(cancellable_manager): } # the record has to outlive the event: the worker still needs to recognize # the signal on its way as a cancellation - cancellable_manager._cache.delete.assert_not_called() + coordination_backend.delete.assert_not_called() -def test_cancel_job_completed_between_read_and_flag(cancellable_manager): +def test_cancel_job_completed_between_read_and_flag( + cancellable_manager, coordination_backend +): """If the job's record is cleared after the auth read, don't revoke.""" - cancellable_manager._cache.get.return_value = json.dumps( + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7} ) # Conditional (xx) write finds no key: the job finished and cleaned up. - cancellable_manager._cache.set.return_value = None + coordination_backend.set.return_value = None with ( mock.patch("superset.extensions.celery_app") as celery_app, @@ -456,8 +457,8 @@ def test_cancel_job_completed_between_read_and_flag(cancellable_manager): celery_app.control.revoke.assert_not_called() -def test_cancel_job_wrong_user_is_rejected(cancellable_manager): - cancellable_manager._cache.get.return_value = json.dumps( +def test_cancel_job_wrong_user_is_rejected(cancellable_manager, coordination_backend): + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7} ) @@ -470,9 +471,11 @@ def test_cancel_job_wrong_user_is_rejected(cancellable_manager): celery_app.control.revoke.assert_not_called() -def test_cancel_job_wrong_channel_is_rejected(cancellable_manager): +def test_cancel_job_wrong_channel_is_rejected( + cancellable_manager, coordination_backend +): """A matching user on a different channel still cannot cancel the job.""" - cancellable_manager._cache.get.return_value = json.dumps( + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7} ) @@ -485,8 +488,8 @@ def test_cancel_job_wrong_channel_is_rejected(cancellable_manager): celery_app.control.revoke.assert_not_called() -def test_cancel_job_unknown_raises(cancellable_manager): - cancellable_manager._cache.get.return_value = None +def test_cancel_job_unknown_raises(cancellable_manager, coordination_backend): + coordination_backend.get.return_value = None with ( mock.patch("superset.extensions.celery_app") as celery_app, @@ -497,39 +500,41 @@ def test_cancel_job_unknown_raises(cancellable_manager): celery_app.control.revoke.assert_not_called() -def test_is_job_cancelled(cancellable_manager): - cancellable_manager._cache.get.return_value = json.dumps( +def test_is_job_cancelled(cancellable_manager, coordination_backend): + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7, "cancelled": True} ) assert cancellable_manager.is_job_cancelled("job-1") is True - cancellable_manager._cache.get.return_value = json.dumps( + coordination_backend.get.return_value = json.dumps( {"channel_id": "chan-1", "user_id": 7} ) assert cancellable_manager.is_job_cancelled("job-1") is False - cancellable_manager._cache.get.return_value = None + coordination_backend.get.return_value = None assert cancellable_manager.is_job_cancelled("job-1") is False -def test_is_job_cancelled_swallows_cache_errors(cancellable_manager): +def test_is_job_cancelled_swallows_cache_errors( + cancellable_manager, coordination_backend +): """A cache failure must not escape and mask the worker's original error.""" - cancellable_manager._cache.get.side_effect = RuntimeError("redis down") + coordination_backend.get.side_effect = RuntimeError("redis down") assert cancellable_manager.is_job_cancelled("job-1") is False -def test_update_job_clears_registry_before_terminal_event(cancellable_manager): +def test_update_job_clears_registry_before_terminal_event( + cancellable_manager, coordination_backend +): """Clearing first is what makes a cancel that lost the race a 404.""" calls = [] cancellable_manager._stream_limit = 100 cancellable_manager._stream_limit_firehose = 1000 - cancellable_manager._cache.delete.side_effect = lambda *_: calls.append("delete") - cancellable_manager._cache.xadd.side_effect = lambda *_: calls.append("xadd") + coordination_backend.delete.side_effect = lambda *_: calls.append("delete") + coordination_backend.xadd.side_effect = lambda *_: calls.append("xadd") job_metadata = {"channel_id": "chan-1", "job_id": "job-1", "user_id": 7} cancellable_manager.update_job(job_metadata, AsyncQueryManager.STATUS_DONE) - cancellable_manager._cache.delete.assert_called_once_with( - "async-events-job-cancel:job-1" - ) + coordination_backend.delete.assert_called_once_with("async-events-job-cancel:job-1") assert calls == ["delete", "xadd", "xadd"] diff --git a/tests/unit_tests/coordination/__init__.py b/tests/unit_tests/coordination/__init__.py new file mode 100644 index 000000000000..13a83393a912 --- /dev/null +++ b/tests/unit_tests/coordination/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py new file mode 100644 index 000000000000..2f18e037e3f5 --- /dev/null +++ b/tests/unit_tests/coordination/test_service.py @@ -0,0 +1,225 @@ +# 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. +from __future__ import annotations + +import threading +from collections.abc import Iterator + +import pytest +from pytest_mock import MockerFixture + +from superset.coordination import ( + CoordinationBackendUnavailableError, + CoordinationService, + SignalListener, +) + + +@pytest.fixture(autouse=True) +def _reset_legacy_state() -> Iterator[None]: + # The legacy backend + warning flag are class-level caches; reset around each + # test so fallback behavior is exercised deterministically. + CoordinationService._legacy_backend = None + CoordinationService._legacy_warning_emitted = False + yield + CoordinationService._legacy_backend = None + CoordinationService._legacy_warning_emitted = False + + +def _patch_distributed_coordination(mocker: MockerFixture, backend: object) -> None: + mocker.patch( + "superset.utils.cache_manager.CacheManager.distributed_coordination", + new_callable=mocker.PropertyMock, + return_value=backend, + ) + + +def test_get_backend_prefers_distributed_coordination( + app_context: None, mocker: MockerFixture +) -> None: + backend = mocker.MagicMock(name="coordination_backend") + _patch_distributed_coordination(mocker, backend) + + assert CoordinationService.get_backend() is backend + assert CoordinationService.is_backend_defined() is True + + +def test_get_backend_none_when_nothing_configured( + app_context: None, mocker: MockerFixture +) -> None: + _patch_distributed_coordination(mocker, None) + mocker.patch.dict( + "flask.current_app.config", {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {}} + ) + + assert CoordinationService.get_backend() is None + assert CoordinationService.is_backend_defined() is False + + +def test_get_backend_falls_back_to_legacy_gaq_backend_with_warning( + app_context: None, mocker: MockerFixture +) -> None: + _patch_distributed_coordination(mocker, None) + mocker.patch.dict( + "flask.current_app.config", + {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {"CACHE_TYPE": "RedisCache"}}, + ) + legacy_backend = mocker.MagicMock(name="legacy_backend") + get_cache_backend = mocker.patch( + "superset.async_events.async_query_manager.get_cache_backend", + return_value=legacy_backend, + ) + warning = mocker.patch("superset.coordination.logger.warning") + + assert CoordinationService.get_backend() is legacy_backend + # The legacy backend is memoized and the deprecation warning emitted once. + assert CoordinationService.get_backend() is legacy_backend + get_cache_backend.assert_called_once() + warning.assert_called_once() + + +def test_backend_only_ops_raise_when_backend_unavailable( + app_context: None, mocker: MockerFixture +) -> None: + _patch_distributed_coordination(mocker, None) + mocker.patch.dict( + "flask.current_app.config", {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {}} + ) + + for op in ( + lambda: CoordinationService.publish("channel", "msg"), + lambda: CoordinationService.get_value("key"), + lambda: CoordinationService.set_value("key", "value"), + lambda: CoordinationService.delete_value("key"), + lambda: CoordinationService.stream_add("stream", {"data": "x"}), + lambda: CoordinationService.stream_range("stream"), + ): + with pytest.raises(CoordinationBackendUnavailableError): + op() + + +def test_ops_delegate_to_backend(app_context: None, mocker: MockerFixture) -> None: + backend = mocker.MagicMock(name="coordination_backend") + backend.publish.return_value = 3 + backend.get.return_value = b"v" + backend.set.return_value = True + backend.delete.return_value = 1 + backend.xadd.return_value = "1-0" + backend.xrange.return_value = [("1-0", {"data": "x"})] + _patch_distributed_coordination(mocker, backend) + + assert CoordinationService.publish("chan", "msg") == 3 + assert CoordinationService.get_value("k") == b"v" + assert CoordinationService.set_value("k", "v", ttl=10, if_present=True) is True + assert CoordinationService.delete_value("k") == 1 + assert CoordinationService.stream_add("stream", {"data": "x"}, "*", 100) == "1-0" + assert CoordinationService.stream_range("stream", "-", "+", 10) == [ + ("1-0", {"data": "x"}) + ] + backend.publish.assert_called_once_with("chan", "msg") + # The generalized set() flags map onto the backend's Redis-native kwargs. + backend.set.assert_called_once_with("k", "v", ex=10, nx=False, xx=True) + backend.delete.assert_called_once_with("k") + backend.xadd.assert_called_once_with("stream", {"data": "x"}, "*", 100) + backend.xrange.assert_called_once_with("stream", "-", "+", 10) + + +# -- wait_for_signal / listen -------------------------------------------------- + + +def test_wait_for_signal_returns_when_check_satisfied( + app_context: None, mocker: MockerFixture +) -> None: + mocker.patch.object(CoordinationService, "get_backend", return_value=None) + assert CoordinationService.wait_for_signal("ch", lambda: "done") == "done" + + +def test_wait_for_signal_polls_until_satisfied_without_backend( + app_context: None, mocker: MockerFixture +) -> None: + mocker.patch.object(CoordinationService, "get_backend", return_value=None) + results = iter([None, "done"]) + + result = CoordinationService.wait_for_signal( + "ch", lambda: next(results), poll_interval=0.01 + ) + assert result == "done" + + +def test_wait_for_signal_times_out(app_context: None, mocker: MockerFixture) -> None: + mocker.patch.object(CoordinationService, "get_backend", return_value=None) + with pytest.raises(TimeoutError): + CoordinationService.wait_for_signal( + "ch", lambda: None, timeout=0.05, poll_interval=0.01 + ) + + +def test_wait_for_signal_wakes_via_pubsub_and_cleans_up( + app_context: None, mocker: MockerFixture +) -> None: + backend = mocker.MagicMock(name="backend") + pubsub = mocker.MagicMock(name="pubsub") + backend.pubsub.return_value = pubsub + mocker.patch.object(CoordinationService, "get_backend", return_value=backend) + results = iter([None, "done"]) + + result = CoordinationService.wait_for_signal( + "ch", lambda: next(results), timeout=5.0 + ) + assert result == "done" + pubsub.subscribe.assert_called_once_with("ch") + # One wake-up nudge between the first (None) and second (done) check. + pubsub.get_message.assert_called_once() + pubsub.unsubscribe.assert_called_once() + pubsub.close.assert_called_once() + + +def test_listen_invokes_on_signal_then_stops( + app_context: None, mocker: MockerFixture +) -> None: + mocker.patch.object(CoordinationService, "get_backend", return_value=None) + fired = threading.Event() + + listener = CoordinationService.listen_for_signal( + "ch", check=lambda: True, on_signal=fired.set, poll_interval=0.01, name="t" + ) + assert fired.wait(timeout=2.0) is True + listener.stop() + + +def test_listen_does_not_fire_when_condition_never_met( + app_context: None, mocker: MockerFixture +) -> None: + mocker.patch.object(CoordinationService, "get_backend", return_value=None) + on_signal = mocker.MagicMock() + + listener = CoordinationService.listen_for_signal( + "ch", check=lambda: False, on_signal=on_signal, poll_interval=0.01 + ) + listener.stop() + on_signal.assert_not_called() + + +def test_signal_listener_stop_signals_and_joins(mocker: MockerFixture) -> None: + thread = mocker.MagicMock(name="thread") + thread.is_alive.side_effect = [True, False] + stop_event = threading.Event() + + SignalListener(thread, stop_event).stop() + + assert stop_event.is_set() + thread.join.assert_called_once_with(timeout=2.0) diff --git a/tests/unit_tests/distributed_lock/distributed_lock_tests.py b/tests/unit_tests/distributed_lock/distributed_lock_tests.py index 3b22c3adc43c..df4980bbd632 100644 --- a/tests/unit_tests/distributed_lock/distributed_lock_tests.py +++ b/tests/unit_tests/distributed_lock/distributed_lock_tests.py @@ -18,16 +18,13 @@ # pylint: disable=invalid-name from typing import Any -from unittest.mock import MagicMock, patch +from unittest.mock import patch from uuid import UUID import pytest from freezegun import freeze_time from sqlalchemy.orm import Session, sessionmaker -# Force module loading before tests run so patches work correctly -import superset.commands.distributed_lock.acquire as acquire_module -import superset.commands.distributed_lock.release as release_module from superset import db from superset.distributed_lock import DistributedLock from superset.distributed_lock.types import LockValue @@ -39,6 +36,12 @@ MAIN_KEY = get_key("ns", a=1, b=2) OTHER_KEY = get_key("ns2", a=1, b=2) +# Distributed locking is plumbed through the coordination service: acquire/release +# call CoordinationService.set_value/delete_value when a backend is defined, else KV. +BACKEND_DEFINED = "superset.coordination.CoordinationService.is_backend_defined" +COORD_SET = "superset.coordination.CoordinationService.set_value" +COORD_DELETE = "superset.coordination.CoordinationService.delete_value" + def _get_lock(key: UUID, session: Session) -> Any: from superset.key_value.models import KeyValueEntry @@ -70,11 +73,8 @@ def test_distributed_lock_kv_happy_path() -> None: """ session = _get_other_session() - # Ensure Redis is not configured so KV backend is used - with ( - patch.object(acquire_module, "get_redis_client", return_value=None), - patch.object(release_module, "get_redis_client", return_value=None), - ): + # Ensure no backend is defined so the KV backend is used + with patch(BACKEND_DEFINED, return_value=False): with freeze_time("2021-01-01"): assert _get_lock(MAIN_KEY, session) is None @@ -100,11 +100,8 @@ def test_distributed_lock_kv_expired() -> None: """ session = _get_other_session() - # Ensure Redis is not configured so KV backend is used - with ( - patch.object(acquire_module, "get_redis_client", return_value=None), - patch.object(release_module, "get_redis_client", return_value=None), - ): + # Ensure no backend is defined so the KV backend is used + with patch(BACKEND_DEFINED, return_value=False): with freeze_time("2021-01-01"): assert _get_lock(MAIN_KEY, session) is None with DistributedLock("ns", a=1, b=2): @@ -116,33 +113,30 @@ def test_distributed_lock_kv_expired() -> None: def test_distributed_lock_uses_redis_when_configured() -> None: - """Test that DistributedLock uses Redis backend when configured.""" - mock_redis = MagicMock() - mock_redis.set.return_value = True # Lock acquired - - # Use patch.object to patch on already-imported modules + """Test that DistributedLock uses the coordination backend when configured.""" with ( - patch.object(acquire_module, "get_redis_client", return_value=mock_redis), - patch.object(release_module, "get_redis_client", return_value=mock_redis), + patch(BACKEND_DEFINED, return_value=True), + patch(COORD_SET, return_value=True) as mock_set, + patch(COORD_DELETE) as mock_delete, ): with DistributedLock("test_redis", key="value") as lock_key: assert lock_key is not None # Verify SET NX EX was called - mock_redis.set.assert_called_once() - call_args = mock_redis.set.call_args + mock_set.assert_called_once() + call_args = mock_set.call_args assert call_args.kwargs["nx"] is True assert "ex" in call_args.kwargs # Verify DELETE was called on exit - mock_redis.delete.assert_called_once() + mock_delete.assert_called_once() def test_distributed_lock_redis_already_taken() -> None: """Test Redis lock fails when already held.""" - mock_redis = MagicMock() - mock_redis.set.return_value = None # Lock not acquired (already taken) - - with patch.object(acquire_module, "get_redis_client", return_value=mock_redis): + with ( + patch(BACKEND_DEFINED, return_value=True), + patch(COORD_SET, return_value=None), # Lock not acquired (already taken) + ): with pytest.raises(AcquireDistributedLockFailedException): with DistributedLock("test_redis", key="value"): pass @@ -152,10 +146,10 @@ def test_distributed_lock_redis_connection_error() -> None: """Test Redis connection error raises exception (fail fast).""" import redis - mock_redis = MagicMock() - mock_redis.set.side_effect = redis.RedisError("Connection failed") - - with patch.object(acquire_module, "get_redis_client", return_value=mock_redis): + with ( + patch(BACKEND_DEFINED, return_value=True), + patch(COORD_SET, side_effect=redis.RedisError("Connection failed")), + ): with pytest.raises(AcquireDistributedLockFailedException): with DistributedLock("test_redis", key="value"): pass @@ -163,15 +157,13 @@ def test_distributed_lock_redis_connection_error() -> None: def test_distributed_lock_custom_ttl() -> None: """Test Redis lock with custom TTL.""" - mock_redis = MagicMock() - mock_redis.set.return_value = True - with ( - patch.object(acquire_module, "get_redis_client", return_value=mock_redis), - patch.object(release_module, "get_redis_client", return_value=mock_redis), + patch(BACKEND_DEFINED, return_value=True), + patch(COORD_SET, return_value=True) as mock_set, + patch(COORD_DELETE), ): with DistributedLock("test", ttl_seconds=60, key="value"): - call_args = mock_redis.set.call_args + call_args = mock_set.call_args assert call_args.kwargs["ex"] == 60 # Custom TTL @@ -179,29 +171,24 @@ def test_distributed_lock_default_ttl(app_context: None) -> None: """Test Redis lock uses default TTL when not specified.""" from superset.commands.distributed_lock.base import get_default_lock_ttl - mock_redis = MagicMock() - mock_redis.set.return_value = True - with ( - patch.object(acquire_module, "get_redis_client", return_value=mock_redis), - patch.object(release_module, "get_redis_client", return_value=mock_redis), + patch(BACKEND_DEFINED, return_value=True), + patch(COORD_SET, return_value=True) as mock_set, + patch(COORD_DELETE), ): with DistributedLock("test", key="value"): - call_args = mock_redis.set.call_args + call_args = mock_set.call_args assert call_args.kwargs["ex"] == get_default_lock_ttl() def test_distributed_lock_fallback_to_kv_when_redis_not_configured() -> None: - """Test falls back to KV lock when Redis not configured.""" + """Test falls back to KV lock when no backend is configured.""" session = _get_other_session() test_key = get_key("test_fallback", key="value") - with ( - patch.object(acquire_module, "get_redis_client", return_value=None), - patch.object(release_module, "get_redis_client", return_value=None), - ): + with patch(BACKEND_DEFINED, return_value=False): with freeze_time("2021-01-01"): - # When Redis is not configured, should use KV backend + # When no backend is defined, should use KV backend with DistributedLock("test_fallback", key="value") as lock_key: assert lock_key == test_key # Verify lock exists in KV store diff --git a/tests/unit_tests/tasks/test_manager.py b/tests/unit_tests/tasks/test_manager.py index 9f10c4da59b4..7690112a4f15 100644 --- a/tests/unit_tests/tasks/test_manager.py +++ b/tests/unit_tests/tasks/test_manager.py @@ -14,71 +14,38 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Unit tests for TaskManager pub/sub functionality""" +"""Unit tests for TaskManager. + +Since PR #43316, the pub/sub-wake-else-poll orchestration lives in +``CoordinationService`` (``publish`` / ``wait_for_signal`` / ``listen_for_signal``); +the tests here cover TaskManager's thin GTF-facing layer: channel naming, publishing +task signals through the service, and delegating waits/abort-listens to it. +""" -import threading -import time from unittest.mock import MagicMock, patch +import pytest import redis -from superset.tasks.manager import AbortListener, TaskManager - - -class TestAbortListener: - """Tests for AbortListener class""" - - def test_stop_sets_event(self): - """Test that stop() sets the stop event""" - stop_event = threading.Event() - thread = MagicMock(spec=threading.Thread) - thread.is_alive.return_value = False +from superset.tasks.manager import TaskManager - listener = AbortListener("test-uuid", thread, stop_event) +GET_BACKEND = "superset.coordination.CoordinationService.get_backend" - assert not stop_event.is_set() - listener.stop() - assert stop_event.is_set() - def test_stop_closes_pubsub(self): - """Test that stop() closes the pub/sub connection""" - stop_event = threading.Event() - thread = MagicMock(spec=threading.Thread) - thread.is_alive.return_value = False - pubsub = MagicMock() - - listener = AbortListener("test-uuid", thread, stop_event, pubsub) - listener.stop() - - pubsub.unsubscribe.assert_called_once() - pubsub.close.assert_called_once() - - def test_stop_joins_thread(self): - """Test that stop() joins the listener thread""" - stop_event = threading.Event() - thread = MagicMock(spec=threading.Thread) - thread.is_alive.return_value = True - - listener = AbortListener("test-uuid", thread, stop_event) - listener.stop() - - thread.join.assert_called_once_with(timeout=2.0) +def _reset_prefixes() -> None: + TaskManager._initialized = False + TaskManager._channel_prefix = "gtf:abort:" + TaskManager._completion_channel_prefix = "gtf:complete:" class TestTaskManagerInitApp: """Tests for TaskManager.init_app()""" def setup_method(self): - """Reset TaskManager state before each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" + _reset_prefixes() def teardown_method(self): - """Reset TaskManager state after each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" + _reset_prefixes() def test_init_app_sets_channel_prefixes(self): """Test init_app reads channel prefixes from config""" @@ -101,357 +68,195 @@ def test_init_app_skips_if_already_initialized(self): app = MagicMock() TaskManager.init_app(app) - # Should not call app.config.get since already initialized app.config.get.assert_not_called() -class TestTaskManagerPubSub: - """Tests for TaskManager pub/sub methods""" +class TestTaskManagerChannels: + """Tests for the abort/completion channel naming.""" def setup_method(self): - """Reset TaskManager state before each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" + _reset_prefixes() def teardown_method(self): - """Reset TaskManager state after each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" - - @patch("superset.tasks.manager.cache_manager") - def test_is_pubsub_available_no_redis(self, mock_cache_manager): - """Test is_pubsub_available returns False when Redis not configured""" - mock_cache_manager.distributed_coordination = None - assert TaskManager.is_pubsub_available() is False - - @patch("superset.tasks.manager.cache_manager") - def test_is_pubsub_available_with_redis(self, mock_cache_manager): - """Test is_pubsub_available returns True when Redis is configured""" - mock_cache_manager.distributed_coordination = MagicMock() - assert TaskManager.is_pubsub_available() is True + _reset_prefixes() def test_get_abort_channel(self): - """Test get_abort_channel returns correct channel name""" - task_uuid = "abc-123-def-456" - channel = TaskManager.get_abort_channel(task_uuid) - assert channel == "gtf:abort:abc-123-def-456" + assert TaskManager.get_abort_channel("abc-123") == "gtf:abort:abc-123" def test_get_abort_channel_custom_prefix(self): - """Test get_abort_channel with custom prefix""" TaskManager._channel_prefix = "custom:prefix:" - task_uuid = "test-uuid" - channel = TaskManager.get_abort_channel(task_uuid) - assert channel == "custom:prefix:test-uuid" - - @patch("superset.tasks.manager.cache_manager") - def test_publish_abort_no_redis(self, mock_cache_manager): - """Test publish_abort returns False when Redis not available""" - mock_cache_manager.distributed_coordination = None - result = TaskManager.publish_abort("test-uuid") - assert result is False - - @patch("superset.tasks.manager.cache_manager") - def test_publish_abort_success(self, mock_cache_manager): - """Test publish_abort publishes message successfully""" - mock_redis = MagicMock() - mock_redis.publish.return_value = 1 # One subscriber - mock_cache_manager.distributed_coordination = mock_redis - - result = TaskManager.publish_abort("test-uuid") + assert TaskManager.get_abort_channel("test-uuid") == "custom:prefix:test-uuid" - assert result is True - mock_redis.publish.assert_called_once_with("gtf:abort:test-uuid", "abort") - - @patch("superset.tasks.manager.cache_manager") - def test_publish_abort_redis_error(self, mock_cache_manager): - """Test publish_abort handles Redis errors gracefully""" - mock_redis = MagicMock() - mock_redis.publish.side_effect = redis.RedisError("Connection lost") - mock_cache_manager.distributed_coordination = mock_redis - - result = TaskManager.publish_abort("test-uuid") + def test_get_completion_channel(self): + assert TaskManager.get_completion_channel("abc-123") == "gtf:complete:abc-123" - assert result is False + def test_get_completion_channel_custom_prefix(self): + TaskManager._completion_channel_prefix = "custom:complete:" + assert ( + TaskManager.get_completion_channel("test-uuid") + == "custom:complete:test-uuid" + ) -class TestTaskManagerListenForAbort: - """Tests for TaskManager.listen_for_abort()""" +class TestTaskManagerPublish: + """publish_abort / publish_completion route through CoordinationService.""" def setup_method(self): - """Reset TaskManager state before each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" + _reset_prefixes() def teardown_method(self): - """Reset TaskManager state after each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" - - @patch("superset.tasks.manager.cache_manager") - def test_listen_for_abort_no_redis_uses_polling(self, mock_cache_manager): - """Test listen_for_abort falls back to polling when Redis unavailable""" - mock_cache_manager.distributed_coordination = None - callback = MagicMock() + _reset_prefixes() - with patch.object(TaskManager, "_poll_for_abort", return_value=None): - listener = TaskManager.listen_for_abort( - task_uuid="test-uuid", - callback=callback, - poll_interval=1.0, - app=None, - ) - - # Give thread time to start - time.sleep(0.1) - listener.stop() + @patch(GET_BACKEND, return_value=None) + def test_publish_abort_no_backend(self, mock_backend): + assert TaskManager.publish_abort("test-uuid") is False - # Should use polling since no Redis - assert listener._pubsub is None + @patch(GET_BACKEND) + def test_publish_abort_success(self, mock_get_backend): + backend = MagicMock() + backend.publish.return_value = 1 + mock_get_backend.return_value = backend - @patch("superset.tasks.manager.cache_manager") - def test_listen_for_abort_with_redis_uses_pubsub(self, mock_cache_manager): - """Test listen_for_abort uses pub/sub when Redis available""" - mock_redis = MagicMock() - mock_pubsub = MagicMock() - mock_redis.pubsub.return_value = mock_pubsub - mock_cache_manager.distributed_coordination = mock_redis + assert TaskManager.publish_abort("test-uuid") is True + backend.publish.assert_called_once_with("gtf:abort:test-uuid", "abort") - callback = MagicMock() + @patch(GET_BACKEND) + def test_publish_abort_redis_error(self, mock_get_backend): + backend = MagicMock() + backend.publish.side_effect = redis.RedisError("Connection lost") + mock_get_backend.return_value = backend - with patch.object(TaskManager, "_listen_pubsub", return_value=None): - listener = TaskManager.listen_for_abort( - task_uuid="test-uuid", - callback=callback, - poll_interval=1.0, - app=None, - ) + assert TaskManager.publish_abort("test-uuid") is False - # Give thread time to start - time.sleep(0.1) - listener.stop() + @patch(GET_BACKEND, return_value=None) + def test_publish_completion_no_backend(self, mock_backend): + assert TaskManager.publish_completion("test-uuid", "success") is False - # Should subscribe to channel - mock_pubsub.subscribe.assert_called_once_with("gtf:abort:test-uuid") + @patch(GET_BACKEND) + def test_publish_completion_success(self, mock_get_backend): + backend = MagicMock() + backend.publish.return_value = 1 + mock_get_backend.return_value = backend - @patch("superset.tasks.manager.cache_manager") - def test_listen_for_abort_redis_subscribe_failure_raises(self, mock_cache_manager): - """Test listen_for_abort raises exception on subscribe failure - when Redis configured""" - import pytest + assert TaskManager.publish_completion("test-uuid", "success") is True + backend.publish.assert_called_once_with("gtf:complete:test-uuid", "success") - mock_redis = MagicMock() - mock_redis.pubsub.side_effect = redis.RedisError("Connection failed") - mock_cache_manager.distributed_coordination = mock_redis + @patch(GET_BACKEND) + def test_publish_completion_redis_error(self, mock_get_backend): + backend = MagicMock() + backend.publish.side_effect = redis.RedisError("Connection lost") + mock_get_backend.return_value = backend - callback = MagicMock() - - # With fail-fast behavior, Redis subscribe failure raises exception - with pytest.raises(redis.RedisError, match="Connection failed"): - TaskManager.listen_for_abort( - task_uuid="test-uuid", - callback=callback, - poll_interval=1.0, - app=None, - ) + assert TaskManager.publish_completion("test-uuid", "success") is False -class TestTaskManagerCompletion: - """Tests for TaskManager completion pub/sub and wait_for_completion""" +class TestTaskManagerListenForAbort: + """listen_for_abort delegates to CoordinationService.listen_for_signal().""" def setup_method(self): - """Reset TaskManager state before each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" + _reset_prefixes() def teardown_method(self): - """Reset TaskManager state after each test""" - TaskManager._initialized = False - TaskManager._channel_prefix = "gtf:abort:" - TaskManager._completion_channel_prefix = "gtf:complete:" + _reset_prefixes() - def test_get_completion_channel(self): - """Test get_completion_channel returns correct channel name""" - task_uuid = "abc-123-def-456" - channel = TaskManager.get_completion_channel(task_uuid) - assert channel == "gtf:complete:abc-123-def-456" + @patch("superset.tasks.manager.TaskManager._check_abort_status") + @patch("superset.coordination.CoordinationService.listen_for_signal") + def test_listen_for_abort_delegates_channel_and_predicate( + self, mock_listen, mock_check + ): + sentinel = MagicMock(name="listener") + mock_listen.return_value = sentinel + callback = MagicMock() - def test_get_completion_channel_custom_prefix(self): - """Test get_completion_channel with custom prefix""" - TaskManager._completion_channel_prefix = "custom:complete:" - task_uuid = "test-uuid" - channel = TaskManager.get_completion_channel(task_uuid) - assert channel == "custom:complete:test-uuid" + listener = TaskManager.listen_for_abort( + task_uuid="test-uuid", + callback=callback, + poll_interval=3.0, + app=None, + ) - @patch("superset.tasks.manager.cache_manager") - def test_publish_completion_no_redis(self, mock_cache_manager): - """Test publish_completion returns False when Redis not available""" - mock_cache_manager.distributed_coordination = None - result = TaskManager.publish_completion("test-uuid", "success") - assert result is False + assert listener is sentinel + args, kwargs = mock_listen.call_args + assert args[0] == "gtf:abort:test-uuid" + assert kwargs["poll_interval"] == 3.0 - @patch("superset.tasks.manager.cache_manager") - def test_publish_completion_success(self, mock_cache_manager): - """Test publish_completion publishes message successfully""" - mock_redis = MagicMock() - mock_redis.publish.return_value = 1 # One subscriber - mock_cache_manager.distributed_coordination = mock_redis + # The predicate closure checks abort status for this task. + kwargs["check"]() + mock_check.assert_called_once_with("test-uuid") - result = TaskManager.publish_completion("test-uuid", "success") + # The on_signal closure invokes the caller's callback. + kwargs["on_signal"]() + callback.assert_called_once() - assert result is True - mock_redis.publish.assert_called_once_with("gtf:complete:test-uuid", "success") - @patch("superset.tasks.manager.cache_manager") - def test_publish_completion_redis_error(self, mock_cache_manager): - """Test publish_completion handles Redis errors gracefully""" - mock_redis = MagicMock() - mock_redis.publish.side_effect = redis.RedisError("Connection lost") - mock_cache_manager.distributed_coordination = mock_redis +class TestTaskManagerWaitForCompletion: + """wait_for_completion delegates the wait to CoordinationService.""" - result = TaskManager.publish_completion("test-uuid", "success") + def setup_method(self): + _reset_prefixes() - assert result is False + def teardown_method(self): + _reset_prefixes() - @patch("superset.tasks.manager.cache_manager") + @patch(GET_BACKEND, return_value=None) @patch("superset.daos.tasks.TaskDAO") - def test_wait_for_completion_task_not_found(self, mock_dao, mock_cache_manager): - """Test wait_for_completion raises ValueError for missing task""" - import pytest - - mock_cache_manager.distributed_coordination = None + def test_task_not_found_raises(self, mock_dao, mock_backend): mock_dao.find_one_or_none.return_value = None - with pytest.raises(ValueError, match="not found"): TaskManager.wait_for_completion("nonexistent-uuid") - @patch("superset.tasks.manager.cache_manager") + @patch(GET_BACKEND, return_value=None) @patch("superset.daos.tasks.TaskDAO") - def test_wait_for_completion_already_complete(self, mock_dao, mock_cache_manager): - """Test wait_for_completion returns immediately for terminal state""" - mock_cache_manager.distributed_coordination = None - mock_task = MagicMock() - mock_task.uuid = "test-uuid" - mock_task.status = "success" - mock_dao.find_one_or_none.return_value = mock_task - - result = TaskManager.wait_for_completion("test-uuid") + def test_already_complete_returns_immediately(self, mock_dao, mock_backend): + task = MagicMock() + task.status = "success" + mock_dao.find_one_or_none.return_value = task - assert result == mock_task - # Should only call find_one_or_none once (initial check) - mock_dao.find_one_or_none.assert_called_once() + assert TaskManager.wait_for_completion("test-uuid") is task - @patch("superset.tasks.manager.cache_manager") + @patch(GET_BACKEND, return_value=None) @patch("superset.daos.tasks.TaskDAO") - def test_wait_for_completion_timeout(self, mock_dao, mock_cache_manager): - """Test wait_for_completion raises TimeoutError when timeout expires""" - import pytest + def test_timeout_raises(self, mock_dao, mock_backend): + task = MagicMock() + task.status = "in_progress" # never terminal + mock_dao.find_one_or_none.return_value = task - mock_cache_manager.distributed_coordination = None - mock_task = MagicMock() - mock_task.uuid = "test-uuid" - mock_task.status = "in_progress" # Never completes - mock_dao.find_one_or_none.return_value = mock_task - - with pytest.raises(TimeoutError, match="Timeout waiting"): - TaskManager.wait_for_completion("test-uuid", timeout=0.1) + with pytest.raises(TimeoutError, match="Timed out waiting"): + TaskManager.wait_for_completion( + "test-uuid", timeout=0.05, poll_interval=0.01 + ) - @patch("superset.tasks.manager.cache_manager") + @patch(GET_BACKEND, return_value=None) @patch("superset.daos.tasks.TaskDAO") - def test_wait_for_completion_polling_success(self, mock_dao, mock_cache_manager): - """Test wait_for_completion returns when task completes via polling""" - mock_cache_manager.distributed_coordination = None - mock_task_pending = MagicMock() - mock_task_pending.uuid = "test-uuid" - mock_task_pending.status = "pending" - - mock_task_complete = MagicMock() - mock_task_complete.uuid = "test-uuid" - mock_task_complete.status = "success" - - # First call returns pending, second returns complete - mock_dao.find_one_or_none.side_effect = [ - mock_task_pending, - mock_task_complete, - ] + def test_polling_success(self, mock_dao, mock_backend): + pending = MagicMock() + pending.status = "pending" + complete = MagicMock() + complete.status = "success" + mock_dao.find_one_or_none.side_effect = [pending, complete] result = TaskManager.wait_for_completion( - "test-uuid", - timeout=5.0, - poll_interval=0.1, + "test-uuid", timeout=5.0, poll_interval=0.01 ) - assert result.status == "success" - @patch("superset.tasks.manager.cache_manager") + @patch(GET_BACKEND) @patch("superset.daos.tasks.TaskDAO") - def test_wait_for_completion_with_pubsub(self, mock_dao, mock_cache_manager): - """Test wait_for_completion uses pub/sub when Redis available""" - mock_task_pending = MagicMock() - mock_task_pending.uuid = "test-uuid" - mock_task_pending.status = "pending" - - mock_task_complete = MagicMock() - mock_task_complete.uuid = "test-uuid" - mock_task_complete.status = "success" - - # First call returns pending, second returns complete - mock_dao.find_one_or_none.side_effect = [ - mock_task_pending, - mock_task_complete, - ] - - # Set up mock Redis with pub/sub - mock_redis = MagicMock() - mock_pubsub = MagicMock() - # Simulate receiving a completion message - mock_pubsub.get_message.return_value = { - "type": "message", - "data": "success", - } - mock_redis.pubsub.return_value = mock_pubsub - mock_cache_manager.distributed_coordination = mock_redis + def test_pubsub_success_subscribes_and_cleans_up(self, mock_dao, mock_get_backend): + pending = MagicMock() + pending.status = "pending" + complete = MagicMock() + complete.status = "success" + mock_dao.find_one_or_none.side_effect = [pending, complete] + + backend = MagicMock() + pubsub = MagicMock() + backend.pubsub.return_value = pubsub + mock_get_backend.return_value = backend - result = TaskManager.wait_for_completion( - "test-uuid", - timeout=5.0, - ) + result = TaskManager.wait_for_completion("test-uuid", timeout=5.0) assert result.status == "success" - # Should have subscribed to completion channel - mock_pubsub.subscribe.assert_called_once_with("gtf:complete:test-uuid") - # Should have cleaned up - mock_pubsub.unsubscribe.assert_called_once() - mock_pubsub.close.assert_called_once() - - @patch("superset.tasks.manager.cache_manager") - @patch("superset.daos.tasks.TaskDAO") - def test_wait_for_completion_pubsub_error_raises( - self, mock_dao, mock_cache_manager - ): - """Test wait_for_completion raises exception on Redis error when - Redis configured""" - import pytest - - mock_task_pending = MagicMock() - mock_task_pending.uuid = "test-uuid" - mock_task_pending.status = "pending" - - mock_dao.find_one_or_none.return_value = mock_task_pending - - # Set up mock Redis that fails - mock_redis = MagicMock() - mock_redis.pubsub.side_effect = redis.RedisError("Connection failed") - mock_cache_manager.distributed_coordination = mock_redis - - # With fail-fast behavior, Redis error is raised instead of falling back - with pytest.raises(redis.RedisError, match="Connection failed"): - TaskManager.wait_for_completion( - "test-uuid", - timeout=5.0, - poll_interval=0.1, - ) + pubsub.subscribe.assert_called_once_with("gtf:complete:test-uuid") + pubsub.unsubscribe.assert_called_once() + pubsub.close.assert_called_once() From aa923831bc6591170bbdad896f9c4970a024d48e Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Wed, 19 Aug 2026 10:06:57 -0700 Subject: [PATCH 02/13] fix(coordination): guard cancel-registry write; repair tests for removed cache_manager and renamed set_value flags --- superset/async_events/async_query_manager.py | 4 ++ .../tasks/test_sync_join_wait.py | 5 +- .../distributed_lock_tests.py | 8 +-- tests/unit_tests/tasks/test_handlers.py | 7 +-- tests/unit_tests/tasks/test_timeout.py | 63 +++++++++---------- 5 files changed, 45 insertions(+), 42 deletions(-) diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index fc357c7aee75..210d2cd427db 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -310,6 +310,10 @@ def _register_cancellable_job( owner without trusting the client-supplied id. Expires with the JWT so it never outlives the job it guards. """ + # Best-effort: the cancel registry is an optimization. Skip when no + # coordination backend is configured rather than failing job submission. + if not CoordinationService.is_backend_defined(): + return CoordinationService.set_value( self._job_registry_key(job_id), json.dumps({"channel_id": channel_id, "user_id": user_id}), diff --git a/tests/integration_tests/tasks/test_sync_join_wait.py b/tests/integration_tests/tasks/test_sync_join_wait.py index 2483c57cd853..c77b313e0db6 100644 --- a/tests/integration_tests/tasks/test_sync_join_wait.py +++ b/tests/integration_tests/tasks/test_sync_join_wait.py @@ -90,8 +90,9 @@ def test_wait_for_completion_timeout(app_context, login_as, get_user) -> None: try: # Force polling mode by mocking distributed_coordination as None - with patch("superset.tasks.manager.cache_manager") as mock_cache_manager: - mock_cache_manager.distributed_coordination = None + with patch( + "superset.coordination.CoordinationService.get_backend", return_value=None + ): with pytest.raises(TimeoutError): TaskManager.wait_for_completion( task.uuid, diff --git a/tests/unit_tests/distributed_lock/distributed_lock_tests.py b/tests/unit_tests/distributed_lock/distributed_lock_tests.py index df4980bbd632..d0bfa000e6b2 100644 --- a/tests/unit_tests/distributed_lock/distributed_lock_tests.py +++ b/tests/unit_tests/distributed_lock/distributed_lock_tests.py @@ -124,8 +124,8 @@ def test_distributed_lock_uses_redis_when_configured() -> None: # Verify SET NX EX was called mock_set.assert_called_once() call_args = mock_set.call_args - assert call_args.kwargs["nx"] is True - assert "ex" in call_args.kwargs + assert call_args.kwargs["if_absent"] is True + assert "ttl" in call_args.kwargs # Verify DELETE was called on exit mock_delete.assert_called_once() @@ -164,7 +164,7 @@ def test_distributed_lock_custom_ttl() -> None: ): with DistributedLock("test", ttl_seconds=60, key="value"): call_args = mock_set.call_args - assert call_args.kwargs["ex"] == 60 # Custom TTL + assert call_args.kwargs["ttl"] == 60 # Custom TTL def test_distributed_lock_default_ttl(app_context: None) -> None: @@ -178,7 +178,7 @@ def test_distributed_lock_default_ttl(app_context: None) -> None: ): with DistributedLock("test", key="value"): call_args = mock_set.call_args - assert call_args.kwargs["ex"] == get_default_lock_ttl() + assert call_args.kwargs["ttl"] == get_default_lock_ttl() def test_distributed_lock_fallback_to_kv_when_redis_not_configured() -> None: diff --git a/tests/unit_tests/tasks/test_handlers.py b/tests/unit_tests/tasks/test_handlers.py index 9e50e82ea7bc..3b30cb0d2ee4 100644 --- a/tests/unit_tests/tasks/test_handlers.py +++ b/tests/unit_tests/tasks/test_handlers.py @@ -80,11 +80,10 @@ def task_context(mock_task, mock_task_dao, mock_update_command, mock_flask_app): with ( patch("superset.tasks.context.current_app") as mock_current_app, - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", return_value=None + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - # Configure current_app mock mock_current_app.config = mock_flask_app.config # Use regular Mock (not MagicMock) for _get_current_object to avoid diff --git a/tests/unit_tests/tasks/test_timeout.py b/tests/unit_tests/tasks/test_timeout.py index f66d55e5a171..41576100fdd9 100644 --- a/tests/unit_tests/tasks/test_timeout.py +++ b/tests/unit_tests/tasks/test_timeout.py @@ -87,11 +87,10 @@ def task_context_for_timeout(mock_flask_app, mock_task_abortable): with ( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", return_value=None + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - # Configure current_app mock mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app @@ -275,11 +274,11 @@ def test_timeout_triggers_abort_when_abortable( patch( "superset.commands.tasks.update.UpdateTaskCommand" ) as mock_update_cmd, - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_abortable @@ -321,11 +320,11 @@ def test_timeout_logs_warning_when_not_abortable( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.tasks.context.logger") as mock_logger, - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_not_abortable @@ -361,11 +360,11 @@ def test_timeout_does_not_trigger_if_already_aborting( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_abortable @@ -467,11 +466,11 @@ def test_timeout_triggered_flag_set_on_timeout( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_abortable @@ -508,11 +507,11 @@ def test_user_abort_does_not_set_timeout_triggered( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_abortable @@ -545,11 +544,11 @@ def test_abort_handlers_completed_tracks_success( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_abortable @@ -582,11 +581,11 @@ def test_abort_handlers_completed_false_on_exception( patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), - patch("superset.tasks.manager.cache_manager") as mock_cache_manager, + patch( + "superset.coordination.CoordinationService.get_backend", + return_value=None, + ), ): - # Disable Redis by making distributed_coordination return None - mock_cache_manager.distributed_coordination = None - mock_current_app.config = mock_flask_app.config mock_current_app._get_current_object.return_value = mock_flask_app mock_dao.find_one_or_none.return_value = mock_task_abortable From e079f8585068b7e6d01b97a36db38e73d7af9db9 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Wed, 19 Aug 2026 11:12:21 -0700 Subject: [PATCH 03/13] fix(coordination): only use the legacy GAQ backend fallback when GLOBAL_ASYNC_QUERIES is enabled MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The deprecated GLOBAL_ASYNC_QUERIES_CACHE_BACKEND ships a populated default, so CoordinationService.get_backend() treated it as a live backend everywhere — making is_backend_defined() return True even with GAQ off and no Redis, and pushing all DistributedLock/GTF callers at a nonexistent Redis. Gate the legacy fallback on the GLOBAL_ASYNC_QUERIES feature flag (the only real signal that GAQ is in use). --- superset/coordination/__init__.py | 11 +++++++++++ .../async_events/async_query_manager_tests.py | 3 ++- tests/unit_tests/coordination/test_service.py | 1 + 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/superset/coordination/__init__.py b/superset/coordination/__init__.py index db6ccb08a438..663d945bbd9d 100644 --- a/superset/coordination/__init__.py +++ b/superset/coordination/__init__.py @@ -145,6 +145,17 @@ def _get_legacy_backend(cls) -> "CoordinationBackend | None": from flask import current_app + from superset import is_feature_enabled + + # GLOBAL_ASYNC_QUERIES_CACHE_BACKEND ships a populated default, so its mere + # presence is not operator intent — it only signals a coordination backend + # when Global Async Queries is actually enabled. Without this gate every + # deployment (and all lock/GTF callers) would treat the default as a live + # Redis backend and try to connect. The legacy bridge exists solely to keep + # GAQ working during the deprecation window. + if not is_feature_enabled("GLOBAL_ASYNC_QUERIES"): + return None + if not current_app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( "CACHE_TYPE" ): diff --git a/tests/unit_tests/async_events/async_query_manager_tests.py b/tests/unit_tests/async_events/async_query_manager_tests.py index f6aaa972ba16..b78e3df1e958 100644 --- a/tests/unit_tests/async_events/async_query_manager_tests.py +++ b/tests/unit_tests/async_events/async_query_manager_tests.py @@ -262,9 +262,10 @@ def test_parse_channel_id_from_request_as_guest_user_differs_per_scope( assert with_datasets != with_rev +@mock.patch("superset.coordination.CoordinationService.get_backend", return_value=None) @mock.patch("superset.is_feature_enabled") def test_submit_chart_data_job_as_guest_user( - is_feature_enabled_mock, async_query_manager + is_feature_enabled_mock, get_backend_mock, async_query_manager ): is_feature_enabled_mock.return_value = True set_current_as_guest_user() diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py index 2f18e037e3f5..88d457ba74e5 100644 --- a/tests/unit_tests/coordination/test_service.py +++ b/tests/unit_tests/coordination/test_service.py @@ -74,6 +74,7 @@ def test_get_backend_falls_back_to_legacy_gaq_backend_with_warning( app_context: None, mocker: MockerFixture ) -> None: _patch_distributed_coordination(mocker, None) + mocker.patch("superset.is_feature_enabled", return_value=True) mocker.patch.dict( "flask.current_app.config", {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {"CACHE_TYPE": "RedisCache"}}, From 41dcdf822e49a769ce305eb3078dd9c267b558e8 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Wed, 19 Aug 2026 12:17:23 -0700 Subject: [PATCH 04/13] refactor(coordination): split types and helpers out of __init__ Move SignalListener + the CoordinationBackend alias to coordination/types.py and the close_pubsub helper to coordination/utils.py, mirroring the superset/distributed_lock layout (types.py + utils.py) and the existing coordination/exceptions.py. __init__ re-exports the public names, so `from superset.coordination import ...` is unchanged. --- superset/coordination/__init__.py | 62 +++----------------------- superset/coordination/types.py | 72 +++++++++++++++++++++++++++++++ superset/coordination/utils.py | 33 ++++++++++++++ 3 files changed, 112 insertions(+), 55 deletions(-) create mode 100644 superset/coordination/types.py create mode 100644 superset/coordination/utils.py diff --git a/superset/coordination/__init__.py b/superset/coordination/__init__.py index 663d945bbd9d..4256af732f55 100644 --- a/superset/coordination/__init__.py +++ b/superset/coordination/__init__.py @@ -41,14 +41,11 @@ from typing import Any, Callable, TYPE_CHECKING, TypeVar from superset.coordination.exceptions import CoordinationBackendUnavailableError +from superset.coordination.types import SignalListener +from superset.coordination.utils import close_pubsub if TYPE_CHECKING: - from superset.async_events.cache_backend import ( - RedisCacheBackend, - RedisSentinelCacheBackend, - ) - - CoordinationBackend = RedisCacheBackend | RedisSentinelCacheBackend + from superset.coordination.types import CoordinationBackend logger = logging.getLogger(__name__) @@ -60,42 +57,6 @@ _PUBSUB_TICK_SECONDS = 1.0 -class SignalListener: - """Handle for a background listener started by - :meth:`CoordinationService.listen_for_signal`. - - Wraps the daemon thread, its stop flag, and (in pub/sub mode) the subscription. - :meth:`stop` sets the flag and closes the subscription so a thread blocked in - ``get_message`` wakes immediately, then joins. - """ - - def __init__( - self, - thread: threading.Thread, - stop_event: threading.Event, - pubsub: Any = None, - ) -> None: - self._thread = thread - self._stop_event = stop_event - self._pubsub = pubsub - - def stop(self) -> None: - """Signal the listener to stop and wait briefly for the thread to finish.""" - self._stop_event.set() - # Closing the subscription unblocks a thread parked in get_message so - # teardown is near-immediate rather than waiting a full poll tick. - if self._pubsub is not None: - _close_pubsub(self._pubsub) - if self._thread.is_alive(): - self._thread.join(timeout=2.0) - if self._thread.is_alive(): - # Daemon thread: it will be reaped at process exit. Don't block. - logger.warning( - "Signal listener thread %s did not terminate within 2s.", - self._thread.name, - ) - - class CoordinationService: """Single entry point for the Valkey/Redis coordination primitives. @@ -335,7 +296,7 @@ def wait_for_signal( cls._wait_tick(pubsub, poll_interval, remaining) finally: if pubsub is not None: - _close_pubsub(pubsub) + close_pubsub(pubsub) @staticmethod def _wait_tick(pubsub: Any, poll_interval: float, remaining: float | None) -> None: @@ -384,7 +345,7 @@ def listen_for_signal( try: pubsub.subscribe(channel) except Exception: - _close_pubsub(pubsub) + close_pubsub(pubsub) raise thread = threading.Thread( target=cls._run_listen_loop, @@ -405,7 +366,7 @@ def _run_listen_loop( poll_interval: float, pubsub: Any, ) -> None: - """Body of the background listener thread (see :meth:`listen`).""" + """Body of the background listener thread (see :meth:`listen_for_signal`).""" try: while not stop_event.is_set(): try: @@ -436,13 +397,4 @@ def _run_listen_loop( logger.exception("Signal listener on %s crashed", channel) finally: if pubsub is not None: - _close_pubsub(pubsub) - - -def _close_pubsub(pubsub: Any) -> None: - """Best-effort unsubscribe + close of a pub/sub subscription.""" - try: - pubsub.unsubscribe() - pubsub.close() - except Exception as ex: # pylint: disable=broad-except - logger.debug("Error closing pub/sub subscription: %s", ex) + close_pubsub(pubsub) diff --git a/superset/coordination/types.py b/superset/coordination/types.py new file mode 100644 index 000000000000..0cb48625c7fa --- /dev/null +++ b/superset/coordination/types.py @@ -0,0 +1,72 @@ +# 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. +"""Types for the coordination service.""" + +from __future__ import annotations + +import logging +import threading +from typing import Any, TYPE_CHECKING + +from superset.coordination.utils import close_pubsub + +if TYPE_CHECKING: + from superset.async_events.cache_backend import ( + RedisCacheBackend, + RedisSentinelCacheBackend, + ) + + # The concrete coordination backend resolved from DISTRIBUTED_COORDINATION_CONFIG. + CoordinationBackend = RedisCacheBackend | RedisSentinelCacheBackend + +logger = logging.getLogger(__name__) + + +class SignalListener: + """Handle for a background listener started by + :meth:`~superset.coordination.CoordinationService.listen_for_signal`. + + Wraps the daemon thread, its stop flag, and (in pub/sub mode) the subscription. + :meth:`stop` sets the flag and closes the subscription so a thread blocked in + ``get_message`` wakes immediately, then joins. + """ + + def __init__( + self, + thread: threading.Thread, + stop_event: threading.Event, + pubsub: Any = None, + ) -> None: + self._thread = thread + self._stop_event = stop_event + self._pubsub = pubsub + + def stop(self) -> None: + """Signal the listener to stop and wait briefly for the thread to finish.""" + self._stop_event.set() + # Closing the subscription unblocks a thread parked in get_message so + # teardown is near-immediate rather than waiting a full poll tick. + if self._pubsub is not None: + close_pubsub(self._pubsub) + if self._thread.is_alive(): + self._thread.join(timeout=2.0) + if self._thread.is_alive(): + # Daemon thread: it will be reaped at process exit. Don't block. + logger.warning( + "Signal listener thread %s did not terminate within 2s.", + self._thread.name, + ) diff --git a/superset/coordination/utils.py b/superset/coordination/utils.py new file mode 100644 index 000000000000..7a567c6b7d31 --- /dev/null +++ b/superset/coordination/utils.py @@ -0,0 +1,33 @@ +# 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. +"""Helpers for the coordination service.""" + +from __future__ import annotations + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + + +def close_pubsub(pubsub: Any) -> None: + """Best-effort unsubscribe + close of a pub/sub subscription.""" + try: + pubsub.unsubscribe() + pubsub.close() + except Exception as ex: # pylint: disable=broad-except + logger.debug("Error closing pub/sub subscription: %s", ex) From fdc28f862d77782c1e6dccdacf3cca7e7b580dd5 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Wed, 19 Aug 2026 15:26:25 -0700 Subject: [PATCH 05/13] refactor(coordination): move CoordinationService to base.py; keep __init__ thin Addresses review feedback: no implementation in __init__.py. CoordinationService now lives in coordination/base.py (matching the base.py convention used elsewhere, e.g. commands/distributed_lock/base.py); __init__ only holds the package docstring and re-exports (CoordinationService, SignalListener, CoordinationBackendUnavailableError), so `from superset.coordination import ...` is unchanged. --- superset/coordination/__init__.py | 368 +--------------------------- superset/coordination/base.py | 386 ++++++++++++++++++++++++++++++ 2 files changed, 392 insertions(+), 362 deletions(-) create mode 100644 superset/coordination/base.py diff --git a/superset/coordination/__init__.py b/superset/coordination/__init__.py index 4256af732f55..19a1d1d451ce 100644 --- a/superset/coordination/__init__.py +++ b/superset/coordination/__init__.py @@ -33,368 +33,12 @@ deprecation warning) so existing deployments keep working during the transition. """ -from __future__ import annotations - -import logging -import threading -import time -from typing import Any, Callable, TYPE_CHECKING, TypeVar - +from superset.coordination.base import CoordinationService from superset.coordination.exceptions import CoordinationBackendUnavailableError from superset.coordination.types import SignalListener -from superset.coordination.utils import close_pubsub - -if TYPE_CHECKING: - from superset.coordination.types import CoordinationBackend - -logger = logging.getLogger(__name__) - -T = TypeVar("T") - -# Poll cadence for the pub/sub wait loop: how long each ``get_message`` blocks -# before the loop re-checks the predicate, the timeout, and the stop flag. Keeps -# stop latency and missed-message recovery bounded to ~1s. -_PUBSUB_TICK_SECONDS = 1.0 - - -class CoordinationService: - """Single entry point for the Valkey/Redis coordination primitives. - - Two layers of API: - - - **Raw primitives** — ``publish``, ``get`` / ``set`` / ``delete``, - ``stream_add`` / ``stream_range``. These are backend-only and have no fallback: - they raise :class:`CoordinationBackendUnavailableError` when no backend is - configured, rather than silently doing nothing. - - **Higher-level await/notify** — ``wait_for_signal`` (blocking) and - ``listen_for_signal`` (background). These combine a pub/sub channel with a - caller-supplied predicate: - when a backend is defined they wake promptly on a published message, and either - way they fall back to polling the predicate. This keeps the pub/sub-vs-poll - boilerplate in one place; callers just supply a channel and a check. - - All methods are class-level: the service is app-global and resolves its backend - from the shared coordination connection on each call. - - Distributed locking is *not* exposed here: it has its own user-facing interface - (:class:`~superset.distributed_lock.DistributedLock`) that uses this service's - backend when one is defined and falls back to a database-backed lock otherwise. - """ - - _legacy_backend: "CoordinationBackend | None" = None - _legacy_warning_emitted: bool = False - - @classmethod - def get_backend(cls) -> "CoordinationBackend | None": - """Resolve the coordination backend. - - Prefers ``DISTRIBUTED_COORDINATION_CONFIG`` (via the cache manager). Falls - back to the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` when only that - is configured, emitting a one-time deprecation warning. Returns ``None`` when - neither is configured. - """ - from superset.extensions import cache_manager - - if (backend := cache_manager.distributed_coordination) is not None: - return backend - return cls._get_legacy_backend() - - @classmethod - def _get_legacy_backend(cls) -> "CoordinationBackend | None": - if cls._legacy_backend is not None: - return cls._legacy_backend - - from flask import current_app - - from superset import is_feature_enabled - - # GLOBAL_ASYNC_QUERIES_CACHE_BACKEND ships a populated default, so its mere - # presence is not operator intent — it only signals a coordination backend - # when Global Async Queries is actually enabled. Without this gate every - # deployment (and all lock/GTF callers) would treat the default as a live - # Redis backend and try to connect. The legacy bridge exists solely to keep - # GAQ working during the deprecation window. - if not is_feature_enabled("GLOBAL_ASYNC_QUERIES"): - return None - - if not current_app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( - "CACHE_TYPE" - ): - return None - - if not cls._legacy_warning_emitted: - logger.warning( - "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND is deprecated and will be " - "removed in Superset 8.0; configure DISTRIBUTED_COORDINATION_CONFIG " - "instead so a single connection powers distributed locks, pub/sub, " - "and streams." - ) - cls._legacy_warning_emitted = True - - from superset.async_events.async_query_manager import get_cache_backend - - cls._legacy_backend = get_cache_backend(current_app.config) - return cls._legacy_backend - - @classmethod - def is_backend_defined(cls) -> bool: - """Whether a coordination backend is defined. - - Some operations require the Valkey/Redis backend - (``DISTRIBUTED_COORDINATION_CONFIG``) to be configured; those that do note it - on their own docstring. Best-effort callers should branch on this before - invoking a backend-dependent operation instead of catching - :class:`CoordinationBackendUnavailableError`. - """ - return cls.get_backend() is not None - - @classmethod - def _require_backend(cls) -> "CoordinationBackend": - """Return the backend or raise if none is configured. - - Used by the backend-only primitives (pub/sub publish, key/value, streams) - so a missing backend fails loudly instead of silently no-op'ing. - """ - backend = cls.get_backend() - if backend is None: - raise CoordinationBackendUnavailableError( - "No coordination backend configured; set " - "DISTRIBUTED_COORDINATION_CONFIG to enable key/value and stream " - "operations." - ) - return backend - - # -- Pub/Sub ------------------------------------------------------------- - - @classmethod - def publish(cls, channel: str, message: str) -> int: - """Publish a message to a channel; returns the subscriber count. - - Only publishing is offered here — subscribing needs the native connection - (a long-lived subscription with its own receive loop), so consumers that - subscribe should obtain it via :meth:`get_backend`. - - :raises CoordinationBackendUnavailableError: if no backend is configured. - """ - return cls._require_backend().publish(channel, message) - - # -- Key/Value ----------------------------------------------------------- - - @classmethod - def get_value(cls, key: str) -> Any: - """Return the raw (bytes) value at ``key``, or ``None`` if absent. - - :raises CoordinationBackendUnavailableError: if no backend is configured. - """ - return cls._require_backend().get(key) - - @classmethod - def set_value( - cls, - key: str, - value: Any, - ttl: int | None = None, - if_absent: bool = False, - if_present: bool = False, - ) -> bool | None: - """Store ``value`` at ``key``. - - :param ttl: optional expiry, in seconds. - :param if_absent: only set if the key does not already exist. - :param if_present: only set if the key already exists. - :returns: ``True`` on success, or ``None`` when an ``if_absent`` / - ``if_present`` condition prevented the write. - :raises CoordinationBackendUnavailableError: if no backend is configured. - """ - return cls._require_backend().set( - key, value, ex=ttl, nx=if_absent, xx=if_present - ) - - @classmethod - def delete_value(cls, *keys: str) -> int: - """Delete one or more keys; returns the number deleted. - - :raises CoordinationBackendUnavailableError: if no backend is configured. - """ - return cls._require_backend().delete(*keys) - - # -- Streams ------------------------------------------------------------- - - @classmethod - def stream_add( - cls, - stream: str, - data: dict[str, Any], - event_id: str = "*", - max_len: int | None = None, - ) -> str: - """Append an event to a stream; returns the generated event id. - - :raises CoordinationBackendUnavailableError: if no backend is configured. - """ - return cls._require_backend().xadd(stream, data, event_id, max_len) - - @classmethod - def stream_range( - cls, - stream: str, - start: str = "-", - end: str = "+", - count: int | None = None, - ) -> list[Any]: - """Read a range of events from a stream. - - :raises CoordinationBackendUnavailableError: if no backend is configured. - """ - return cls._require_backend().xrange(stream, start, end, count) - - # -- Await / notify ------------------------------------------------------ - - @classmethod - def wait_for_signal( - cls, - channel: str, - check: Callable[[], T | None], - *, - timeout: float | None = None, - poll_interval: float = 1.0, - ) -> T: - """Block until ``check()`` returns a non-``None`` value; return that value. - - ``check`` is the source of truth (typically a metastore read). When a - coordination backend is defined, this subscribes to ``channel`` and re-runs - ``check`` promptly whenever a message is published; otherwise it polls - ``check`` every ``poll_interval`` seconds. ``check`` is also re-evaluated on - every tick even in pub/sub mode, so a signal published before the subscription - (or a dropped message) is still caught. - - :param channel: pub/sub channel that peers publish to when the awaited state - is reached (used only as a low-latency wake-up; correctness relies on - ``check``). - :param check: returns a truthy result once the wait is satisfied, else - ``None``. - :param timeout: max seconds to wait; ``None`` waits indefinitely. - :param poll_interval: poll cadence when no backend is defined. - :raises TimeoutError: if ``timeout`` elapses before ``check`` is satisfied. - """ - deadline = None if timeout is None else time.monotonic() + timeout - backend = cls.get_backend() - pubsub = backend.pubsub() if backend is not None else None - try: - if pubsub is not None: - pubsub.subscribe(channel) - while True: - # ``check`` is the source of truth; run it first so the fast path and - # any signal missed before subscribing are both covered. - if (result := check()) is not None: - return result - remaining = ( - None if deadline is None else max(0.0, deadline - time.monotonic()) - ) - if remaining is not None and remaining <= 0: - raise TimeoutError(f"Timed out waiting on channel {channel}") - cls._wait_tick(pubsub, poll_interval, remaining) - finally: - if pubsub is not None: - close_pubsub(pubsub) - - @staticmethod - def _wait_tick(pubsub: Any, poll_interval: float, remaining: float | None) -> None: - """Block for one wait tick: a pub/sub message (nudge) or a poll sleep.""" - if pubsub is not None: - wait = ( - _PUBSUB_TICK_SECONDS - if remaining is None - else min(_PUBSUB_TICK_SECONDS, remaining) - ) - pubsub.get_message(ignore_subscribe_messages=True, timeout=wait) - else: - time.sleep( - poll_interval if remaining is None else min(poll_interval, remaining) - ) - - @classmethod - def listen_for_signal( - cls, - channel: str, - check: Callable[[], bool], - on_signal: Callable[[], None], - *, - poll_interval: float, - name: str | None = None, - ) -> SignalListener: - """Run a background daemon that invokes ``on_signal`` once ``check`` is true. - - Same wake-vs-poll model as :meth:`wait_for_signal`: a published message on - ``channel`` wakes the loop when a backend is defined, otherwise it polls - ``check`` every ``poll_interval`` seconds. The thread stops after firing - ``on_signal`` once, or when :meth:`SignalListener.stop` is called. - - :param channel: pub/sub channel peers publish to when the condition is met. - :param check: returns ``True`` once ``on_signal`` should fire. - :param on_signal: invoked (once) when ``check`` becomes true. - :param poll_interval: poll cadence when no backend is defined. - :param name: optional thread name suffix for logging. - """ - stop_event = threading.Event() - backend = cls.get_backend() - pubsub = backend.pubsub() if backend is not None else None - if pubsub is not None: - # Subscribe in the caller's thread so a connection failure surfaces here - # (fail-fast) rather than dying silently in the daemon thread. - try: - pubsub.subscribe(channel) - except Exception: - close_pubsub(pubsub) - raise - thread = threading.Thread( - target=cls._run_listen_loop, - args=(channel, check, on_signal, stop_event, poll_interval, pubsub), - daemon=True, - name=f"coord-listen-{name or channel}", - ) - thread.start() - return SignalListener(thread, stop_event, pubsub) - @classmethod - def _run_listen_loop( - cls, - channel: str, - check: Callable[[], bool], - on_signal: Callable[[], None], - stop_event: threading.Event, - poll_interval: float, - pubsub: Any, - ) -> None: - """Body of the background listener thread (see :meth:`listen_for_signal`).""" - try: - while not stop_event.is_set(): - try: - if check(): - on_signal() - return - if pubsub is not None: - # Blocks up to a tick; the message is just a wake-up nudge. - pubsub.get_message( - ignore_subscribe_messages=True, - timeout=_PUBSUB_TICK_SECONDS, - ) - else: - stop_event.wait(timeout=poll_interval) - except (ValueError, OSError) as ex: - # Connection torn down (e.g. stop() closing the subscription, or - # shutdown). Expected when stopping; otherwise surface it and bail. - if not stop_event.is_set(): - logger.error( - "Signal listener on %s failed: %s", - channel, - ex, - exc_info=True, - ) - return - except Exception: # pylint: disable=broad-except - if not stop_event.is_set(): - logger.exception("Signal listener on %s crashed", channel) - finally: - if pubsub is not None: - close_pubsub(pubsub) +__all__ = [ + "CoordinationBackendUnavailableError", + "CoordinationService", + "SignalListener", +] diff --git a/superset/coordination/base.py b/superset/coordination/base.py new file mode 100644 index 000000000000..7d599ba50b23 --- /dev/null +++ b/superset/coordination/base.py @@ -0,0 +1,386 @@ +# 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. +"""Coordination service implementation. + +See :mod:`superset.coordination` for the package overview. +""" + +from __future__ import annotations + +import logging +import threading +import time +from typing import Any, Callable, TYPE_CHECKING, TypeVar + +from superset.coordination.exceptions import CoordinationBackendUnavailableError +from superset.coordination.types import SignalListener +from superset.coordination.utils import close_pubsub + +if TYPE_CHECKING: + from superset.coordination.types import CoordinationBackend + +logger = logging.getLogger(__name__) + +T = TypeVar("T") + +# Poll cadence for the pub/sub wait loop: how long each ``get_message`` blocks +# before the loop re-checks the predicate, the timeout, and the stop flag. Keeps +# stop latency and missed-message recovery bounded to ~1s. +_PUBSUB_TICK_SECONDS = 1.0 + + +class CoordinationService: + """Single entry point for the Valkey/Redis coordination primitives. + + Two layers of API: + + - **Raw primitives** — ``publish``, ``get`` / ``set`` / ``delete``, + ``stream_add`` / ``stream_range``. These are backend-only and have no fallback: + they raise :class:`CoordinationBackendUnavailableError` when no backend is + configured, rather than silently doing nothing. + - **Higher-level await/notify** — ``wait_for_signal`` (blocking) and + ``listen_for_signal`` (background). These combine a pub/sub channel with a + caller-supplied predicate: + when a backend is defined they wake promptly on a published message, and either + way they fall back to polling the predicate. This keeps the pub/sub-vs-poll + boilerplate in one place; callers just supply a channel and a check. + + All methods are class-level: the service is app-global and resolves its backend + from the shared coordination connection on each call. + + Distributed locking is *not* exposed here: it has its own user-facing interface + (:class:`~superset.distributed_lock.DistributedLock`) that uses this service's + backend when one is defined and falls back to a database-backed lock otherwise. + """ + + _legacy_backend: "CoordinationBackend | None" = None + _legacy_warning_emitted: bool = False + + @classmethod + def get_backend(cls) -> "CoordinationBackend | None": + """Resolve the coordination backend. + + Prefers ``DISTRIBUTED_COORDINATION_CONFIG`` (via the cache manager). Falls + back to the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` when only that + is configured, emitting a one-time deprecation warning. Returns ``None`` when + neither is configured. + """ + from superset.extensions import cache_manager + + if (backend := cache_manager.distributed_coordination) is not None: + return backend + return cls._get_legacy_backend() + + @classmethod + def _get_legacy_backend(cls) -> "CoordinationBackend | None": + if cls._legacy_backend is not None: + return cls._legacy_backend + + from flask import current_app + + from superset import is_feature_enabled + + # GLOBAL_ASYNC_QUERIES_CACHE_BACKEND ships a populated default, so its mere + # presence is not operator intent — it only signals a coordination backend + # when Global Async Queries is actually enabled. Without this gate every + # deployment (and all lock/GTF callers) would treat the default as a live + # Redis backend and try to connect. The legacy bridge exists solely to keep + # GAQ working during the deprecation window. + if not is_feature_enabled("GLOBAL_ASYNC_QUERIES"): + return None + + if not current_app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( + "CACHE_TYPE" + ): + return None + + if not cls._legacy_warning_emitted: + logger.warning( + "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND is deprecated and will be " + "removed in Superset 8.0; configure DISTRIBUTED_COORDINATION_CONFIG " + "instead so a single connection powers distributed locks, pub/sub, " + "and streams." + ) + cls._legacy_warning_emitted = True + + from superset.async_events.async_query_manager import get_cache_backend + + cls._legacy_backend = get_cache_backend(current_app.config) + return cls._legacy_backend + + @classmethod + def is_backend_defined(cls) -> bool: + """Whether a coordination backend is defined. + + Some operations require the Valkey/Redis backend + (``DISTRIBUTED_COORDINATION_CONFIG``) to be configured; those that do note it + on their own docstring. Best-effort callers should branch on this before + invoking a backend-dependent operation instead of catching + :class:`CoordinationBackendUnavailableError`. + """ + return cls.get_backend() is not None + + @classmethod + def _require_backend(cls) -> "CoordinationBackend": + """Return the backend or raise if none is configured. + + Used by the backend-only primitives (pub/sub publish, key/value, streams) + so a missing backend fails loudly instead of silently no-op'ing. + """ + backend = cls.get_backend() + if backend is None: + raise CoordinationBackendUnavailableError( + "No coordination backend configured; set " + "DISTRIBUTED_COORDINATION_CONFIG to enable key/value and stream " + "operations." + ) + return backend + + # -- Pub/Sub ------------------------------------------------------------- + + @classmethod + def publish(cls, channel: str, message: str) -> int: + """Publish a message to a channel; returns the subscriber count. + + Only publishing is offered here — subscribing needs the native connection + (a long-lived subscription with its own receive loop), so consumers that + subscribe should obtain it via :meth:`get_backend`. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().publish(channel, message) + + # -- Key/Value ----------------------------------------------------------- + + @classmethod + def get_value(cls, key: str) -> Any: + """Return the raw (bytes) value at ``key``, or ``None`` if absent. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().get(key) + + @classmethod + def set_value( + cls, + key: str, + value: Any, + ttl: int | None = None, + if_absent: bool = False, + if_present: bool = False, + ) -> bool | None: + """Store ``value`` at ``key``. + + :param ttl: optional expiry, in seconds. + :param if_absent: only set if the key does not already exist. + :param if_present: only set if the key already exists. + :returns: ``True`` on success, or ``None`` when an ``if_absent`` / + ``if_present`` condition prevented the write. + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().set( + key, value, ex=ttl, nx=if_absent, xx=if_present + ) + + @classmethod + def delete_value(cls, *keys: str) -> int: + """Delete one or more keys; returns the number deleted. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().delete(*keys) + + # -- Streams ------------------------------------------------------------- + + @classmethod + def stream_add( + cls, + stream: str, + data: dict[str, Any], + event_id: str = "*", + max_len: int | None = None, + ) -> str: + """Append an event to a stream; returns the generated event id. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().xadd(stream, data, event_id, max_len) + + @classmethod + def stream_range( + cls, + stream: str, + start: str = "-", + end: str = "+", + count: int | None = None, + ) -> list[Any]: + """Read a range of events from a stream. + + :raises CoordinationBackendUnavailableError: if no backend is configured. + """ + return cls._require_backend().xrange(stream, start, end, count) + + # -- Await / notify ------------------------------------------------------ + + @classmethod + def wait_for_signal( + cls, + channel: str, + check: Callable[[], T | None], + *, + timeout: float | None = None, + poll_interval: float = 1.0, + ) -> T: + """Block until ``check()`` returns a non-``None`` value; return that value. + + ``check`` is the source of truth (typically a metastore read). When a + coordination backend is defined, this subscribes to ``channel`` and re-runs + ``check`` promptly whenever a message is published; otherwise it polls + ``check`` every ``poll_interval`` seconds. ``check`` is also re-evaluated on + every tick even in pub/sub mode, so a signal published before the subscription + (or a dropped message) is still caught. + + :param channel: pub/sub channel that peers publish to when the awaited state + is reached (used only as a low-latency wake-up; correctness relies on + ``check``). + :param check: returns a truthy result once the wait is satisfied, else + ``None``. + :param timeout: max seconds to wait; ``None`` waits indefinitely. + :param poll_interval: poll cadence when no backend is defined. + :raises TimeoutError: if ``timeout`` elapses before ``check`` is satisfied. + """ + deadline = None if timeout is None else time.monotonic() + timeout + backend = cls.get_backend() + pubsub = backend.pubsub() if backend is not None else None + try: + if pubsub is not None: + pubsub.subscribe(channel) + while True: + # ``check`` is the source of truth; run it first so the fast path and + # any signal missed before subscribing are both covered. + if (result := check()) is not None: + return result + remaining = ( + None if deadline is None else max(0.0, deadline - time.monotonic()) + ) + if remaining is not None and remaining <= 0: + raise TimeoutError(f"Timed out waiting on channel {channel}") + cls._wait_tick(pubsub, poll_interval, remaining) + finally: + if pubsub is not None: + close_pubsub(pubsub) + + @staticmethod + def _wait_tick(pubsub: Any, poll_interval: float, remaining: float | None) -> None: + """Block for one wait tick: a pub/sub message (nudge) or a poll sleep.""" + if pubsub is not None: + wait = ( + _PUBSUB_TICK_SECONDS + if remaining is None + else min(_PUBSUB_TICK_SECONDS, remaining) + ) + pubsub.get_message(ignore_subscribe_messages=True, timeout=wait) + else: + time.sleep( + poll_interval if remaining is None else min(poll_interval, remaining) + ) + + @classmethod + def listen_for_signal( + cls, + channel: str, + check: Callable[[], bool], + on_signal: Callable[[], None], + *, + poll_interval: float, + name: str | None = None, + ) -> SignalListener: + """Run a background daemon that invokes ``on_signal`` once ``check`` is true. + + Same wake-vs-poll model as :meth:`wait_for_signal`: a published message on + ``channel`` wakes the loop when a backend is defined, otherwise it polls + ``check`` every ``poll_interval`` seconds. The thread stops after firing + ``on_signal`` once, or when :meth:`SignalListener.stop` is called. + + :param channel: pub/sub channel peers publish to when the condition is met. + :param check: returns ``True`` once ``on_signal`` should fire. + :param on_signal: invoked (once) when ``check`` becomes true. + :param poll_interval: poll cadence when no backend is defined. + :param name: optional thread name suffix for logging. + """ + stop_event = threading.Event() + backend = cls.get_backend() + pubsub = backend.pubsub() if backend is not None else None + if pubsub is not None: + # Subscribe in the caller's thread so a connection failure surfaces here + # (fail-fast) rather than dying silently in the daemon thread. + try: + pubsub.subscribe(channel) + except Exception: + close_pubsub(pubsub) + raise + thread = threading.Thread( + target=cls._run_listen_loop, + args=(channel, check, on_signal, stop_event, poll_interval, pubsub), + daemon=True, + name=f"coord-listen-{name or channel}", + ) + thread.start() + return SignalListener(thread, stop_event, pubsub) + + @classmethod + def _run_listen_loop( + cls, + channel: str, + check: Callable[[], bool], + on_signal: Callable[[], None], + stop_event: threading.Event, + poll_interval: float, + pubsub: Any, + ) -> None: + """Body of the background listener thread (see :meth:`listen_for_signal`).""" + try: + while not stop_event.is_set(): + try: + if check(): + on_signal() + return + if pubsub is not None: + # Blocks up to a tick; the message is just a wake-up nudge. + pubsub.get_message( + ignore_subscribe_messages=True, + timeout=_PUBSUB_TICK_SECONDS, + ) + else: + stop_event.wait(timeout=poll_interval) + except (ValueError, OSError) as ex: + # Connection torn down (e.g. stop() closing the subscription, or + # shutdown). Expected when stopping; otherwise surface it and bail. + if not stop_event.is_set(): + logger.error( + "Signal listener on %s failed: %s", + channel, + ex, + exc_info=True, + ) + return + except Exception: # pylint: disable=broad-except + if not stop_event.is_set(): + logger.exception("Signal listener on %s crashed", channel) + finally: + if pubsub is not None: + close_pubsub(pubsub) From 20f9bab2d0f9606df504f43cb80a63026cc6a156 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Wed, 19 Aug 2026 16:17:35 -0700 Subject: [PATCH 06/13] refactor(coordination): import from submodules, empty __init__ Drop the re-exports and __all__ from superset/coordination/__init__.py so the package init carries only the docstring, per review feedback that we don't use the __init__ re-export pattern here. All consumers now import directly from the submodule that owns the symbol: - CoordinationService -> superset.coordination.base - SignalListener -> superset.coordination.types - CoordinationBackendUnavailableError -> superset.coordination.exceptions Test patch targets and docstring cross-refs updated to the .base. paths. --- superset/async_events/async_query_manager.py | 2 +- superset/commands/distributed_lock/acquire.py | 2 +- superset/commands/distributed_lock/release.py | 2 +- superset/config.py | 2 +- superset/coordination/__init__.py | 15 +++++---------- superset/coordination/exceptions.py | 2 +- superset/coordination/types.py | 2 +- superset/tasks/context.py | 2 +- superset/tasks/manager.py | 10 +++++----- .../integration_tests/async_events/api_tests.py | 2 +- .../tasks/test_sync_join_wait.py | 3 ++- .../async_events/async_query_manager_tests.py | 6 ++++-- tests/unit_tests/coordination/test_service.py | 10 ++++------ .../distributed_lock/distributed_lock_tests.py | 6 +++--- tests/unit_tests/tasks/test_handlers.py | 3 ++- tests/unit_tests/tasks/test_manager.py | 4 ++-- tests/unit_tests/tasks/test_timeout.py | 17 +++++++++-------- 17 files changed, 44 insertions(+), 46 deletions(-) diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index 210d2cd427db..7e272ef6e2e0 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -30,7 +30,7 @@ RedisCacheBackend, RedisSentinelCacheBackend, ) -from superset.coordination import CoordinationService +from superset.coordination.base import CoordinationService from superset.utils import json from superset.utils.core import get_user_id diff --git a/superset/commands/distributed_lock/acquire.py b/superset/commands/distributed_lock/acquire.py index d12a979d8feb..901dca4b7c41 100644 --- a/superset/commands/distributed_lock/acquire.py +++ b/superset/commands/distributed_lock/acquire.py @@ -28,7 +28,7 @@ BaseDistributedLockCommand, get_default_lock_ttl, ) -from superset.coordination import CoordinationService +from superset.coordination.base import CoordinationService from superset.daos.key_value import KeyValueDAO from superset.exceptions import ( AcquireDistributedLockFailedException, diff --git a/superset/commands/distributed_lock/release.py b/superset/commands/distributed_lock/release.py index 65221a731a5c..42fdfe0e4c11 100644 --- a/superset/commands/distributed_lock/release.py +++ b/superset/commands/distributed_lock/release.py @@ -24,7 +24,7 @@ from sqlalchemy.exc import SQLAlchemyError from superset.commands.distributed_lock.base import BaseDistributedLockCommand -from superset.coordination import CoordinationService +from superset.coordination.base import CoordinationService from superset.daos.key_value import KeyValueDAO from superset.exceptions import ReleaseDistributedLockFailedException from superset.key_value.exceptions import KeyValueDeleteFailedError diff --git a/superset/config.py b/superset/config.py index 46c7b3d0c35a..a4b6afa71ca6 100644 --- a/superset/config.py +++ b/superset/config.py @@ -3239,7 +3239,7 @@ class ExtraAccessQueryFilters(TypedDict, total=False): # - Global Async Queries event streams (the async-events / firehose transport) # # This backend powers the higher-level coordination service -# (``superset.coordination.CoordinationService``) exposing standardized interfaces +# (``superset.coordination.base.CoordinationService``) exposing standardized interfaces # for distributed locks, pub/sub, and streams under a single connection. Global # Async Queries use this when configured; the former # ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` is deprecated and, if set, used only as a diff --git a/superset/coordination/__init__.py b/superset/coordination/__init__.py index 19a1d1d451ce..9fb3958661af 100644 --- a/superset/coordination/__init__.py +++ b/superset/coordination/__init__.py @@ -31,14 +31,9 @@ reusable coordination surface, and reduces the number of moving parts. The legacy ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` is still honored as a fallback (with a deprecation warning) so existing deployments keep working during the transition. -""" - -from superset.coordination.base import CoordinationService -from superset.coordination.exceptions import CoordinationBackendUnavailableError -from superset.coordination.types import SignalListener -__all__ = [ - "CoordinationBackendUnavailableError", - "CoordinationService", - "SignalListener", -] +Import concrete classes directly from their modules: +:class:`~superset.coordination.base.CoordinationService`, +:class:`~superset.coordination.types.SignalListener`, and +:class:`~superset.coordination.exceptions.CoordinationBackendUnavailableError`. +""" diff --git a/superset/coordination/exceptions.py b/superset/coordination/exceptions.py index 00e2b794afdf..17f9214c5458 100644 --- a/superset/coordination/exceptions.py +++ b/superset/coordination/exceptions.py @@ -26,6 +26,6 @@ class CoordinationBackendUnavailableError(Exception): calling them without a configured coordination backend is a programming or configuration error rather than a silently-ignored no-op. Callers that have their own fallback (e.g. database polling) should gate on - :meth:`superset.coordination.CoordinationService.is_backend_defined` instead of + :meth:`superset.coordination.base.CoordinationService.is_backend_defined` instead of catching this. """ diff --git a/superset/coordination/types.py b/superset/coordination/types.py index 0cb48625c7fa..fa73f586dfef 100644 --- a/superset/coordination/types.py +++ b/superset/coordination/types.py @@ -38,7 +38,7 @@ class SignalListener: """Handle for a background listener started by - :meth:`~superset.coordination.CoordinationService.listen_for_signal`. + :meth:`~superset.coordination.base.CoordinationService.listen_for_signal`. Wraps the daemon thread, its stop flag, and (in pub/sub mode) the subscription. :meth:`stop` sets the flag and closes the subscription so a thread blocked in diff --git a/superset/tasks/context.py b/superset/tasks/context.py index b4052022e558..048ab7a76fc2 100644 --- a/superset/tasks/context.py +++ b/superset/tasks/context.py @@ -35,7 +35,7 @@ from superset.tasks.utils import progress_update if TYPE_CHECKING: - from superset.coordination import SignalListener + from superset.coordination.types import SignalListener from superset.models.tasks import Task logger = logging.getLogger(__name__) diff --git a/superset/tasks/manager.py b/superset/tasks/manager.py index b1ab21c98da7..564df0ac2713 100644 --- a/superset/tasks/manager.py +++ b/superset/tasks/manager.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: from flask import Flask - from superset.coordination import SignalListener + from superset.coordination.types import SignalListener from superset.models.tasks import Task logger = logging.getLogger(__name__) @@ -94,7 +94,7 @@ def publish_abort(cls, task_uuid: UUID) -> bool: :param task_uuid: UUID of the task to abort :returns: True if message was published, False if Redis unavailable """ - from superset.coordination import CoordinationService + from superset.coordination.base import CoordinationService if not CoordinationService.is_backend_defined(): return False @@ -134,7 +134,7 @@ def publish_completion(cls, task_uuid: UUID, status: str) -> bool: :param status: Final status of the task :returns: True if message was published, False if Redis unavailable """ - from superset.coordination import CoordinationService + from superset.coordination.base import CoordinationService if not CoordinationService.is_backend_defined(): return False @@ -177,7 +177,7 @@ def wait_for_completion( :raises TimeoutError: If timeout expires before task completes :raises ValueError: If task not found """ - from superset.coordination import CoordinationService + from superset.coordination.base import CoordinationService from superset.daos.tasks import TaskDAO def get_task() -> "Task | None": @@ -227,7 +227,7 @@ def listen_for_abort( :param app: Flask app for database access in background thread :returns: SignalListener handle to stop listening """ - from superset.coordination import CoordinationService + from superset.coordination.base import CoordinationService def in_context(fn: Callable[[], Any]) -> Callable[[], Any]: # The listener runs in a background thread; DB access needs app context. diff --git a/tests/integration_tests/async_events/api_tests.py b/tests/integration_tests/async_events/api_tests.py index bceca0845fbe..3968acb19ee2 100644 --- a/tests/integration_tests/async_events/api_tests.py +++ b/tests/integration_tests/async_events/api_tests.py @@ -54,7 +54,7 @@ def run_test_with_cache_backend(self, cache_backend_cls: Type[Any], test_func): self.login(ADMIN_USERNAME) with mock.patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=mock_cache, ): test_func(mock_cache) diff --git a/tests/integration_tests/tasks/test_sync_join_wait.py b/tests/integration_tests/tasks/test_sync_join_wait.py index c77b313e0db6..7fb07c4cf690 100644 --- a/tests/integration_tests/tasks/test_sync_join_wait.py +++ b/tests/integration_tests/tasks/test_sync_join_wait.py @@ -91,7 +91,8 @@ def test_wait_for_completion_timeout(app_context, login_as, get_user) -> None: try: # Force polling mode by mocking distributed_coordination as None with patch( - "superset.coordination.CoordinationService.get_backend", return_value=None + "superset.coordination.base.CoordinationService.get_backend", + return_value=None, ): with pytest.raises(TimeoutError): TaskManager.wait_for_completion( diff --git a/tests/unit_tests/async_events/async_query_manager_tests.py b/tests/unit_tests/async_events/async_query_manager_tests.py index b78e3df1e958..c95e0e150f94 100644 --- a/tests/unit_tests/async_events/async_query_manager_tests.py +++ b/tests/unit_tests/async_events/async_query_manager_tests.py @@ -262,7 +262,9 @@ def test_parse_channel_id_from_request_as_guest_user_differs_per_scope( assert with_datasets != with_rev -@mock.patch("superset.coordination.CoordinationService.get_backend", return_value=None) +@mock.patch( + "superset.coordination.base.CoordinationService.get_backend", return_value=None +) @mock.patch("superset.is_feature_enabled") def test_submit_chart_data_job_as_guest_user( is_feature_enabled_mock, get_backend_mock, async_query_manager @@ -364,7 +366,7 @@ def coordination_backend(): """Patch the coordination service to a mock Redis backend and expose it.""" backend = mock.Mock(spec=RedisCacheBackend) with mock.patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=backend, ): yield backend diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py index 88d457ba74e5..6fe869ba15c4 100644 --- a/tests/unit_tests/coordination/test_service.py +++ b/tests/unit_tests/coordination/test_service.py @@ -22,11 +22,9 @@ import pytest from pytest_mock import MockerFixture -from superset.coordination import ( - CoordinationBackendUnavailableError, - CoordinationService, - SignalListener, -) +from superset.coordination.base import CoordinationService +from superset.coordination.exceptions import CoordinationBackendUnavailableError +from superset.coordination.types import SignalListener @pytest.fixture(autouse=True) @@ -84,7 +82,7 @@ def test_get_backend_falls_back_to_legacy_gaq_backend_with_warning( "superset.async_events.async_query_manager.get_cache_backend", return_value=legacy_backend, ) - warning = mocker.patch("superset.coordination.logger.warning") + warning = mocker.patch("superset.coordination.base.logger.warning") assert CoordinationService.get_backend() is legacy_backend # The legacy backend is memoized and the deprecation warning emitted once. diff --git a/tests/unit_tests/distributed_lock/distributed_lock_tests.py b/tests/unit_tests/distributed_lock/distributed_lock_tests.py index d0bfa000e6b2..b9eb4b4d136a 100644 --- a/tests/unit_tests/distributed_lock/distributed_lock_tests.py +++ b/tests/unit_tests/distributed_lock/distributed_lock_tests.py @@ -38,9 +38,9 @@ # Distributed locking is plumbed through the coordination service: acquire/release # call CoordinationService.set_value/delete_value when a backend is defined, else KV. -BACKEND_DEFINED = "superset.coordination.CoordinationService.is_backend_defined" -COORD_SET = "superset.coordination.CoordinationService.set_value" -COORD_DELETE = "superset.coordination.CoordinationService.delete_value" +BACKEND_DEFINED = "superset.coordination.base.CoordinationService.is_backend_defined" +COORD_SET = "superset.coordination.base.CoordinationService.set_value" +COORD_DELETE = "superset.coordination.base.CoordinationService.delete_value" def _get_lock(key: UUID, session: Session) -> Any: diff --git a/tests/unit_tests/tasks/test_handlers.py b/tests/unit_tests/tasks/test_handlers.py index 3b30cb0d2ee4..0105c55ce96b 100644 --- a/tests/unit_tests/tasks/test_handlers.py +++ b/tests/unit_tests/tasks/test_handlers.py @@ -81,7 +81,8 @@ def task_context(mock_task, mock_task_dao, mock_update_command, mock_flask_app): with ( patch("superset.tasks.context.current_app") as mock_current_app, patch( - "superset.coordination.CoordinationService.get_backend", return_value=None + "superset.coordination.base.CoordinationService.get_backend", + return_value=None, ), ): # Configure current_app mock diff --git a/tests/unit_tests/tasks/test_manager.py b/tests/unit_tests/tasks/test_manager.py index 7690112a4f15..96afe46fd6d2 100644 --- a/tests/unit_tests/tasks/test_manager.py +++ b/tests/unit_tests/tasks/test_manager.py @@ -29,7 +29,7 @@ from superset.tasks.manager import TaskManager -GET_BACKEND = "superset.coordination.CoordinationService.get_backend" +GET_BACKEND = "superset.coordination.base.CoordinationService.get_backend" def _reset_prefixes() -> None: @@ -160,7 +160,7 @@ def teardown_method(self): _reset_prefixes() @patch("superset.tasks.manager.TaskManager._check_abort_status") - @patch("superset.coordination.CoordinationService.listen_for_signal") + @patch("superset.coordination.base.CoordinationService.listen_for_signal") def test_listen_for_abort_delegates_channel_and_predicate( self, mock_listen, mock_check ): diff --git a/tests/unit_tests/tasks/test_timeout.py b/tests/unit_tests/tasks/test_timeout.py index 41576100fdd9..073fd1878da3 100644 --- a/tests/unit_tests/tasks/test_timeout.py +++ b/tests/unit_tests/tasks/test_timeout.py @@ -88,7 +88,8 @@ def task_context_for_timeout(mock_flask_app, mock_task_abortable): patch("superset.tasks.context.current_app") as mock_current_app, patch("superset.daos.tasks.TaskDAO") as mock_dao, patch( - "superset.coordination.CoordinationService.get_backend", return_value=None + "superset.coordination.base.CoordinationService.get_backend", + return_value=None, ), ): # Configure current_app mock @@ -275,7 +276,7 @@ def test_timeout_triggers_abort_when_abortable( "superset.commands.tasks.update.UpdateTaskCommand" ) as mock_update_cmd, patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): @@ -321,7 +322,7 @@ def test_timeout_logs_warning_when_not_abortable( patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.tasks.context.logger") as mock_logger, patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): @@ -361,7 +362,7 @@ def test_timeout_does_not_trigger_if_already_aborting( patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): @@ -467,7 +468,7 @@ def test_timeout_triggered_flag_set_on_timeout( patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): @@ -508,7 +509,7 @@ def test_user_abort_does_not_set_timeout_triggered( patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): @@ -545,7 +546,7 @@ def test_abort_handlers_completed_tracks_success( patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): @@ -582,7 +583,7 @@ def test_abort_handlers_completed_false_on_exception( patch("superset.daos.tasks.TaskDAO") as mock_dao, patch("superset.commands.tasks.update.UpdateTaskCommand"), patch( - "superset.coordination.CoordinationService.get_backend", + "superset.coordination.base.CoordinationService.get_backend", return_value=None, ), ): From 947ce35d9fc8ab3fdc4e15ec95efbd8c10e96acb Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Thu, 20 Aug 2026 17:12:59 -0700 Subject: [PATCH 07/13] refactor(coordination): scope legacy GAQ backend to GAQ; fix wait_for_signal fast path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback on #43316. get_backend() now resolves DISTRIBUTED_COORDINATION_CONFIG only. Previously it fell back to the deprecated GLOBAL_ASYNC_QUERIES_CACHE_BACKEND for *all* callers, so a deployment with GAQ enabled but no DISTRIBUTED_COORDINATION_CONFIG silently moved distributed locks (DB -> GAQ Redis) and GTF (DB polling -> GAQ Redis pub/sub) onto the GAQ backend — a rolling-upgrade split-brain for non-GAQ consumers (OAuth2/thumbnails/reports locks included). Locks and GTF are now byte-for-byte master behavior (DB when no coordinator is configured). Global Async Queries keep their own separate backend during the deprecation window: AsyncQueryManager resolves it (get_cache_backend, preferring the dedicated GLOBAL_ASYNC_QUERIES_CACHE_BACKEND and otherwise the coordinator) and passes it to the CoordinationService primitives via a new optional backend= param. The dual-backend arrangement is deprecated and collapses to DISTRIBUTED_COORDINATION_CONFIG in 8.0 (config.py + UPDATING.md). wait_for_signal now runs check() once before opening the pub/sub subscription and returns immediately if satisfied, so an already-terminal task never requires Redis to be reachable (and skips an unnecessary round-trip). --- UPDATING.md | 2 +- superset/async_events/async_query_manager.py | 67 +++++++-- superset/config.py | 24 ++-- superset/coordination/base.py | 132 ++++++++---------- .../async_events/api_tests.py | 8 +- .../async_events/async_query_manager_tests.py | 17 +-- tests/unit_tests/coordination/test_service.py | 70 ++++++---- 7 files changed, 184 insertions(+), 136 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index dba1f305c559..614f31ea98a0 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -61,7 +61,7 @@ the old counter to use the outcome-specific replacements. - [42393](https://github.com/apache/superset/pull/42393): Exported dataset YAML now carries a `uuid` for each metric and column so that custom folder assignments (which reference metrics/columns by UUID) survive an import into another workspace. This affects any export bundle that contains datasets, not just a dataset export: chart, dashboard, database and full-asset exports all embed the same dataset YAML, so a dashboard exported from this release also fails to import into an older one even though no dataset was exported directly. As with `folders` and `currency_code_column`, the affected `datasets/` files fail schema validation (`Unknown field: uuid`) when imported into Superset releases that predate this change; regenerate or hand-edit exports for older targets in mixed-version fleets. - [42300](https://github.com/apache/superset/pull/42300): Timeseries charts (line/area/bar) with a Y-axis bound in effect — either an explicit `yAxisBounds` or one derived from `truncateYAxis` — now clamp out-of-range data points to that bound instead of letting ECharts drop the point (and the line segments around it) entirely. Any existing chart with a configured Y-axis bound and data outside it will look different after upgrading: a gap becomes a point pinned to the boundary. The clamp also rewrites the value ECharts reads for that point's tooltip and data label, so the displayed value is the bound rather than the true observation. - [42087](https://github.com/apache/superset/pull/42087): Stored calculated-column and metric expressions are validated when a query is built, under the same sub-query policy already applied to adhoc expressions. Previously only the dataset update path checked them on save, so expressions written by v1 import, by dataset duplication, or before that check existed were never validated. Since `ALLOW_ADHOC_SUBQUERY` defaults to `False` (see [19242](https://github.com/apache/superset/pull/19242)), a dataset whose stored expression contains a sub-query works before upgrading and afterwards fails at chart render with `Custom SQL fields cannot contain sub-queries.` There is no migration step, and the error does not name the offending dataset column, so audit stored expressions before upgrading: either rewrite them without the sub-query, or set `ALLOW_ADHOC_SUBQUERY = True` to keep the previous behaviour for both stored and adhoc expressions. -- The `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` config key is **deprecated** in favor of the shared coordination backend `DISTRIBUTED_COORDINATION_CONFIG`, which powers a single Redis connection for distributed locks, pub/sub, and the async-events streams (the GAQ firehose). All parameters previously accepted by `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` are supported identically under `DISTRIBUTED_COORDINATION_CONFIG` (they use the same `RedisCache`/`RedisSentinelCache` backend). Existing configs continue to work: when `DISTRIBUTED_COORDINATION_CONFIG` is unset, `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` is used as a fallback and logs a one-time deprecation warning. Migrate by renaming the key; the deprecated key will be removed in Superset 8.0. +- The `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` config key is **deprecated** in favor of the shared coordination backend `DISTRIBUTED_COORDINATION_CONFIG`, which powers a single Redis connection for distributed locks, pub/sub, and the async-events streams (the GAQ firehose). All parameters previously accepted by `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` are supported identically under `DISTRIBUTED_COORDINATION_CONFIG` (they use the same `RedisCache`/`RedisSentinelCache` backend). During the deprecation window the two backends stay scoped: distributed locks and the Global Task Framework use `DISTRIBUTED_COORDINATION_CONFIG` exclusively (falling back to the metadata database when it is unset, unchanged from before), and Global Async Queries keep running on their own `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` (falling back to `DISTRIBUTED_COORDINATION_CONFIG` when that dedicated backend is unset), logging a one-time deprecation warning. This dual-backend arrangement is removed in Superset 8.0, when GAQ moves onto `DISTRIBUTED_COORDINATION_CONFIG`; migrate now by configuring `DISTRIBUTED_COORDINATION_CONFIG`. ### Selenium support removed — Playwright is now required for screenshots diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index 7e272ef6e2e0..95bcc09ffa38 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -89,8 +89,9 @@ def get_cache_backend( ) -> RedisCacheBackend | RedisSentinelCacheBackend: """Build a coordination backend from the deprecated GAQ cache config. - DEPRECATED: retained only so the legacy ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` - setting keeps working as a fallback. Prefer ``DISTRIBUTED_COORDINATION_CONFIG``. + DEPRECATED: retained only so Global Async Queries can run on its own dedicated + ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` during the deprecation window. Removed in + Superset 8.0, when GAQ moves onto ``DISTRIBUTED_COORDINATION_CONFIG``. """ cache_config = config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}) cache_type = cache_config.get("CACHE_TYPE") @@ -114,12 +115,18 @@ class AsyncQueryManager: # Redis key prefix (within the GAQ stream namespace) for the per-job record # that authorizes cancellation and flags a job as cancelled for the worker. _JOB_REGISTRY_PREFIX = "job-cancel:" + # Emit the dedicated-backend deprecation notice at most once per process. + _legacy_backend_warning_emitted: bool = False def __init__(self) -> None: super().__init__() self._stream_prefix: str = "" self._stream_limit: Optional[int] self._stream_limit_firehose: Optional[int] + # Global Async Queries owns its coordination backend separately from the + # shared coordinator (see init_app); resolved there and passed explicitly + # to CoordinationService's primitives. + self._gaq_backend: RedisCacheBackend | RedisSentinelCacheBackend | None = None self._jwt_cookie_name: str = "" self._jwt_cookie_secure: bool = False self._jwt_cookie_domain: Optional[str] @@ -152,6 +159,28 @@ def init_app(self, app: Flask) -> None: "is deprecated)." ) + # Global Async Queries keeps its own coordination backend during the + # deprecation window: prefer the dedicated (deprecated) + # GLOBAL_ASYNC_QUERIES_CACHE_BACKEND when configured, otherwise use the + # shared DISTRIBUTED_COORDINATION_CONFIG. This scopes GAQ's stream/pub-sub + # traffic to its own connection and keeps it off the coordinator's backend + # (which powers distributed locks and the Global Task Framework). In 8.0 the + # dedicated backend is removed and GAQ moves onto the coordinator's connection. + if app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get("CACHE_TYPE"): + if not AsyncQueryManager._legacy_backend_warning_emitted: + logger.warning( + "Global Async Queries is running on its own " + "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND. This dedicated backend is " + "deprecated and will be removed in Superset 8.0, when GAQ will " + "use DISTRIBUTED_COORDINATION_CONFIG like the rest of Superset's " + "coordination (distributed locks, task framework, pub/sub). " + "Configure DISTRIBUTED_COORDINATION_CONFIG to consolidate now." + ) + AsyncQueryManager._legacy_backend_warning_emitted = True + self._gaq_backend = get_cache_backend(app.config) + else: + self._gaq_backend = CoordinationService.get_backend() + if len(app.config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"]) < 32: raise AsyncQueryTokenException( "Please provide a JWT secret at least 32 bytes long" @@ -312,12 +341,13 @@ def _register_cancellable_job( """ # Best-effort: the cancel registry is an optimization. Skip when no # coordination backend is configured rather than failing job submission. - if not CoordinationService.is_backend_defined(): + if self._gaq_backend is None: return CoordinationService.set_value( self._job_registry_key(job_id), json.dumps({"channel_id": channel_id, "user_id": user_id}), ttl=self._jwt_expiration_seconds or None, + backend=self._gaq_backend, ) def submit_chart_data_job( @@ -351,13 +381,13 @@ def submit_chart_data_job( def read_events( self, channel: str, last_id: Optional[str] ) -> list[Optional[dict[str, Any]]]: - if not CoordinationService.is_backend_defined(): + if self._gaq_backend is None: raise CacheBackendNotInitialized("Cache backend not initialized") stream_name = f"{self._stream_prefix}{channel}" start_id = increment_id(last_id) if last_id else "-" results = CoordinationService.stream_range( - stream_name, start_id, "+", self.MAX_EVENT_COUNT + stream_name, start_id, "+", self.MAX_EVENT_COUNT, backend=self._gaq_backend ) # Decode bytes to strings: the coordination Redis backends do not enable # decode_responses, so stream_range returns raw bytes. @@ -376,7 +406,7 @@ def read_events( def update_job( self, job_metadata: dict[str, Any], status: str, **kwargs: Any ) -> None: - if not CoordinationService.is_backend_defined(): + if self._gaq_backend is None: raise CacheBackendNotInitialized("Cache backend not initialized") if "channel_id" not in job_metadata: @@ -401,13 +431,23 @@ def update_job( # SIGUSR1 it is about to receive as a cancellation. if status in (self.STATUS_DONE, self.STATUS_ERROR): if job_id := job_metadata.get("job_id"): - CoordinationService.delete_value(self._job_registry_key(job_id)) + CoordinationService.delete_value( + self._job_registry_key(job_id), backend=self._gaq_backend + ) CoordinationService.stream_add( - scoped_stream_name, event_data, "*", self._stream_limit + scoped_stream_name, + event_data, + "*", + self._stream_limit, + backend=self._gaq_backend, ) CoordinationService.stream_add( - full_stream_name, event_data, "*", self._stream_limit_firehose + full_stream_name, + event_data, + "*", + self._stream_limit_firehose, + backend=self._gaq_backend, ) def is_job_cancelled(self, job_id: str) -> bool: @@ -419,7 +459,9 @@ def is_job_cancelled(self, job_id: str) -> bool: the original error (e.g. a genuine timeout) with a connection error. """ try: - raw = CoordinationService.get_value(self._job_registry_key(job_id)) + raw = CoordinationService.get_value( + self._job_registry_key(job_id), backend=self._gaq_backend + ) if raw is None: return False return bool(json.loads(raw).get("cancelled")) @@ -443,11 +485,11 @@ def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int]) -> No :raises AsyncQueryJobException: the job is unknown or already terminal :raises AsyncQueryTokenException: the caller does not own the job """ - if not CoordinationService.is_backend_defined(): + if self._gaq_backend is None: raise CacheBackendNotInitialized("Cache backend not initialized") key = self._job_registry_key(job_id) - raw = CoordinationService.get_value(key) + raw = CoordinationService.get_value(key, backend=self._gaq_backend) if raw is None: raise AsyncQueryJobException("Job not found or already completed") @@ -465,6 +507,7 @@ def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int]) -> No json.dumps({**record, "cancelled": True}), ttl=self._jwt_expiration_seconds or None, if_present=True, + backend=self._gaq_backend, ) if not flagged: raise AsyncQueryJobException("Job not found or already completed") diff --git a/superset/config.py b/superset/config.py index a4b6afa71ca6..b6b3d814e56c 100644 --- a/superset/config.py +++ b/superset/config.py @@ -2942,11 +2942,15 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument # noq # - Set 'CACHE_TYPE' to 'RedisCache' for RedisCacheBackend. # - Set 'CACHE_TYPE' to 'RedisSentinelCache' for RedisSentinelCacheBackend. # -# DEPRECATED: prefer DISTRIBUTED_COORDINATION_CONFIG, which powers a single -# coordination service for distributed locks, pub/sub, and the async-events -# streams. When DISTRIBUTED_COORDINATION_CONFIG is set it takes precedence; this -# setting is only used as a fallback (with a deprecation warning) and will be -# removed in Superset 8.0. All parameters here are supported identically under +# DEPRECATED: this dedicated backend is retained only so Global Async Queries can +# keep running on its own coordination connection during the deprecation window. +# When configured it is used by GAQ *only* (its event streams and cancel registry); +# it never powers distributed locks or the Global Task Framework, which use +# DISTRIBUTED_COORDINATION_CONFIG exclusively. GAQ prefers this dedicated backend +# when set and otherwise falls back to DISTRIBUTED_COORDINATION_CONFIG. This +# dual-backend arrangement is deprecated and removed in Superset 8.0, when GAQ +# moves onto DISTRIBUTED_COORDINATION_CONFIG like the rest of Superset's +# coordination. All parameters here are supported identically under # DISTRIBUTED_COORDINATION_CONFIG (both use the same # RedisCache/RedisSentinelCache backend). GLOBAL_ASYNC_QUERIES_CACHE_BACKEND = { @@ -3240,10 +3244,12 @@ class ExtraAccessQueryFilters(TypedDict, total=False): # # This backend powers the higher-level coordination service # (``superset.coordination.base.CoordinationService``) exposing standardized interfaces -# for distributed locks, pub/sub, and streams under a single connection. Global -# Async Queries use this when configured; the former -# ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` is deprecated and, if set, used only as a -# fallback (with a deprecation warning). +# for distributed locks, pub/sub, and streams under a single connection. It is the +# single source of truth for the coordinator's consumers (distributed locks, the +# Global Task Framework, and future stream/pub-sub users). Global Async Queries keep +# their own dedicated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` for now (falling back to +# this when that is unset); that dual-backend arrangement is deprecated and, in +# Superset 8.0, GAQ moves onto this connection and the dedicated backend is removed. # # All parameters previously supported by GLOBAL_ASYNC_QUERIES_CACHE_BACKEND are # supported here (both go through the same RedisCacheBackend/RedisSentinelCacheBackend diff --git a/superset/coordination/base.py b/superset/coordination/base.py index 7d599ba50b23..2f8961b76692 100644 --- a/superset/coordination/base.py +++ b/superset/coordination/base.py @@ -51,7 +51,9 @@ class CoordinationService: - **Raw primitives** — ``publish``, ``get`` / ``set`` / ``delete``, ``stream_add`` / ``stream_range``. These are backend-only and have no fallback: they raise :class:`CoordinationBackendUnavailableError` when no backend is - configured, rather than silently doing nothing. + available, rather than silently doing nothing. Each accepts an optional + ``backend`` so a caller with its own connection (Global Async Queries, during + the deprecation window) can run against it instead of the shared coordinator. - **Higher-level await/notify** — ``wait_for_signal`` (blocking) and ``listen_for_signal`` (background). These combine a pub/sub channel with a caller-supplied predicate: @@ -67,60 +69,22 @@ class CoordinationService: backend when one is defined and falls back to a database-backed lock otherwise. """ - _legacy_backend: "CoordinationBackend | None" = None - _legacy_warning_emitted: bool = False - @classmethod def get_backend(cls) -> "CoordinationBackend | None": - """Resolve the coordination backend. - - Prefers ``DISTRIBUTED_COORDINATION_CONFIG`` (via the cache manager). Falls - back to the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` when only that - is configured, emitting a one-time deprecation warning. Returns ``None`` when - neither is configured. + """Resolve the coordination backend from ``DISTRIBUTED_COORDINATION_CONFIG``. + + Returns the shared coordination connection (via the cache manager), or + ``None`` when ``DISTRIBUTED_COORDINATION_CONFIG`` is not configured. This is + the single source of truth for the coordinator's consumers (distributed + locks, the Global Task Framework, and future stream/pub-sub users); it does + *not* consult the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND``. Global + Async Queries owns its own separate backend during the deprecation window — + see :class:`~superset.async_events.async_query_manager.AsyncQueryManager` — + and passes it explicitly to the primitives below via ``backend``. """ from superset.extensions import cache_manager - if (backend := cache_manager.distributed_coordination) is not None: - return backend - return cls._get_legacy_backend() - - @classmethod - def _get_legacy_backend(cls) -> "CoordinationBackend | None": - if cls._legacy_backend is not None: - return cls._legacy_backend - - from flask import current_app - - from superset import is_feature_enabled - - # GLOBAL_ASYNC_QUERIES_CACHE_BACKEND ships a populated default, so its mere - # presence is not operator intent — it only signals a coordination backend - # when Global Async Queries is actually enabled. Without this gate every - # deployment (and all lock/GTF callers) would treat the default as a live - # Redis backend and try to connect. The legacy bridge exists solely to keep - # GAQ working during the deprecation window. - if not is_feature_enabled("GLOBAL_ASYNC_QUERIES"): - return None - - if not current_app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( - "CACHE_TYPE" - ): - return None - - if not cls._legacy_warning_emitted: - logger.warning( - "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND is deprecated and will be " - "removed in Superset 8.0; configure DISTRIBUTED_COORDINATION_CONFIG " - "instead so a single connection powers distributed locks, pub/sub, " - "and streams." - ) - cls._legacy_warning_emitted = True - - from superset.async_events.async_query_manager import get_cache_backend - - cls._legacy_backend = get_cache_backend(current_app.config) - return cls._legacy_backend + return cache_manager.distributed_coordination @classmethod def is_backend_defined(cls) -> bool: @@ -135,13 +99,18 @@ def is_backend_defined(cls) -> bool: return cls.get_backend() is not None @classmethod - def _require_backend(cls) -> "CoordinationBackend": - """Return the backend or raise if none is configured. + def _require_backend( + cls, backend: "CoordinationBackend | None" = None + ) -> "CoordinationBackend": + """Return a usable backend or raise if none is available. Used by the backend-only primitives (pub/sub publish, key/value, streams) - so a missing backend fails loudly instead of silently no-op'ing. + so a missing backend fails loudly instead of silently no-op'ing. When + ``backend`` is supplied (e.g. Global Async Queries passing its own separate + backend) it is used directly; otherwise the shared coordinator backend is + resolved via :meth:`get_backend`. """ - backend = cls.get_backend() + backend = backend or cls.get_backend() if backend is None: raise CoordinationBackendUnavailableError( "No coordination backend configured; set " @@ -153,26 +122,33 @@ def _require_backend(cls) -> "CoordinationBackend": # -- Pub/Sub ------------------------------------------------------------- @classmethod - def publish(cls, channel: str, message: str) -> int: + def publish( + cls, + channel: str, + message: str, + backend: "CoordinationBackend | None" = None, + ) -> int: """Publish a message to a channel; returns the subscriber count. Only publishing is offered here — subscribing needs the native connection (a long-lived subscription with its own receive loop), so consumers that subscribe should obtain it via :meth:`get_backend`. - :raises CoordinationBackendUnavailableError: if no backend is configured. + :param backend: optional explicit backend (see :meth:`_require_backend`). + :raises CoordinationBackendUnavailableError: if no backend is available. """ - return cls._require_backend().publish(channel, message) + return cls._require_backend(backend).publish(channel, message) # -- Key/Value ----------------------------------------------------------- @classmethod - def get_value(cls, key: str) -> Any: + def get_value(cls, key: str, backend: "CoordinationBackend | None" = None) -> Any: """Return the raw (bytes) value at ``key``, or ``None`` if absent. - :raises CoordinationBackendUnavailableError: if no backend is configured. + :param backend: optional explicit backend (see :meth:`_require_backend`). + :raises CoordinationBackendUnavailableError: if no backend is available. """ - return cls._require_backend().get(key) + return cls._require_backend(backend).get(key) @classmethod def set_value( @@ -182,27 +158,32 @@ def set_value( ttl: int | None = None, if_absent: bool = False, if_present: bool = False, + backend: "CoordinationBackend | None" = None, ) -> bool | None: """Store ``value`` at ``key``. :param ttl: optional expiry, in seconds. :param if_absent: only set if the key does not already exist. :param if_present: only set if the key already exists. + :param backend: optional explicit backend (see :meth:`_require_backend`). :returns: ``True`` on success, or ``None`` when an ``if_absent`` / ``if_present`` condition prevented the write. - :raises CoordinationBackendUnavailableError: if no backend is configured. + :raises CoordinationBackendUnavailableError: if no backend is available. """ - return cls._require_backend().set( + return cls._require_backend(backend).set( key, value, ex=ttl, nx=if_absent, xx=if_present ) @classmethod - def delete_value(cls, *keys: str) -> int: + def delete_value( + cls, *keys: str, backend: "CoordinationBackend | None" = None + ) -> int: """Delete one or more keys; returns the number deleted. - :raises CoordinationBackendUnavailableError: if no backend is configured. + :param backend: optional explicit backend (see :meth:`_require_backend`). + :raises CoordinationBackendUnavailableError: if no backend is available. """ - return cls._require_backend().delete(*keys) + return cls._require_backend(backend).delete(*keys) # -- Streams ------------------------------------------------------------- @@ -213,12 +194,14 @@ def stream_add( data: dict[str, Any], event_id: str = "*", max_len: int | None = None, + backend: "CoordinationBackend | None" = None, ) -> str: """Append an event to a stream; returns the generated event id. - :raises CoordinationBackendUnavailableError: if no backend is configured. + :param backend: optional explicit backend (see :meth:`_require_backend`). + :raises CoordinationBackendUnavailableError: if no backend is available. """ - return cls._require_backend().xadd(stream, data, event_id, max_len) + return cls._require_backend(backend).xadd(stream, data, event_id, max_len) @classmethod def stream_range( @@ -227,12 +210,14 @@ def stream_range( start: str = "-", end: str = "+", count: int | None = None, + backend: "CoordinationBackend | None" = None, ) -> list[Any]: """Read a range of events from a stream. - :raises CoordinationBackendUnavailableError: if no backend is configured. + :param backend: optional explicit backend (see :meth:`_require_backend`). + :raises CoordinationBackendUnavailableError: if no backend is available. """ - return cls._require_backend().xrange(stream, start, end, count) + return cls._require_backend(backend).xrange(stream, start, end, count) # -- Await / notify ------------------------------------------------------ @@ -264,14 +249,19 @@ def wait_for_signal( :raises TimeoutError: if ``timeout`` elapses before ``check`` is satisfied. """ deadline = None if timeout is None else time.monotonic() + timeout + # Check first, before touching the backend: if the awaited state is already + # reached (e.g. the task is already terminal), return straight from the + # source of truth so the fast path never requires the backend to be reachable. + if (result := check()) is not None: + return result backend = cls.get_backend() pubsub = backend.pubsub() if backend is not None else None try: if pubsub is not None: pubsub.subscribe(channel) while True: - # ``check`` is the source of truth; run it first so the fast path and - # any signal missed before subscribing are both covered. + # Re-check every tick even in pub/sub mode, so a signal published + # before the subscription (or a dropped message) is still caught. if (result := check()) is not None: return result remaining = ( diff --git a/tests/integration_tests/async_events/api_tests.py b/tests/integration_tests/async_events/api_tests.py index 3968acb19ee2..c711453c550c 100644 --- a/tests/integration_tests/async_events/api_tests.py +++ b/tests/integration_tests/async_events/api_tests.py @@ -46,17 +46,17 @@ def cancel_event(self, job_id: str): def run_test_with_cache_backend(self, cache_backend_cls: Type[Any], test_func): app._got_first_request = False - async_query_manager_factory.init_app(app) - # The manager resolves its backend through CoordinationService.get_backend(), - # so patch there rather than assigning a (now-removed) private attribute. + # GAQ resolves its own backend from get_cache_backend during init_app, so + # inject the mock there and initialize within the patch. mock_cache = mock.Mock(spec=cache_backend_cls) self.login(ADMIN_USERNAME) with mock.patch( - "superset.coordination.base.CoordinationService.get_backend", + "superset.async_events.async_query_manager.get_cache_backend", return_value=mock_cache, ): + async_query_manager_factory.init_app(app) test_func(mock_cache) def _test_events_logic(self, mock_cache): diff --git a/tests/unit_tests/async_events/async_query_manager_tests.py b/tests/unit_tests/async_events/async_query_manager_tests.py index c95e0e150f94..91dbd267e056 100644 --- a/tests/unit_tests/async_events/async_query_manager_tests.py +++ b/tests/unit_tests/async_events/async_query_manager_tests.py @@ -262,16 +262,15 @@ def test_parse_channel_id_from_request_as_guest_user_differs_per_scope( assert with_datasets != with_rev -@mock.patch( - "superset.coordination.base.CoordinationService.get_backend", return_value=None -) @mock.patch("superset.is_feature_enabled") def test_submit_chart_data_job_as_guest_user( - is_feature_enabled_mock, get_backend_mock, async_query_manager + is_feature_enabled_mock, async_query_manager ): is_feature_enabled_mock.return_value = True set_current_as_guest_user() + # The manager has no GAQ backend wired (init_app not run), so the best-effort + # cancel registry write is skipped and submission still proceeds. job_mock = Mock() async_query_manager._load_chart_data_into_cache_job = job_mock job_meta = async_query_manager.submit_chart_data_job( @@ -363,13 +362,8 @@ def test_view(): @fixture def coordination_backend(): - """Patch the coordination service to a mock Redis backend and expose it.""" - backend = mock.Mock(spec=RedisCacheBackend) - with mock.patch( - "superset.coordination.base.CoordinationService.get_backend", - return_value=backend, - ): - yield backend + """A mock Redis backend GAQ operates against as its own dedicated backend.""" + return mock.Mock(spec=RedisCacheBackend) @fixture @@ -378,6 +372,7 @@ def cancellable_manager(coordination_backend): manager = AsyncQueryManager() manager._jwt_expiration_seconds = 3600 manager._stream_prefix = "async-events-" + manager._gaq_backend = coordination_backend return manager diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py index 6fe869ba15c4..49250492c11a 100644 --- a/tests/unit_tests/coordination/test_service.py +++ b/tests/unit_tests/coordination/test_service.py @@ -17,7 +17,6 @@ from __future__ import annotations import threading -from collections.abc import Iterator import pytest from pytest_mock import MockerFixture @@ -27,17 +26,6 @@ from superset.coordination.types import SignalListener -@pytest.fixture(autouse=True) -def _reset_legacy_state() -> Iterator[None]: - # The legacy backend + warning flag are class-level caches; reset around each - # test so fallback behavior is exercised deterministically. - CoordinationService._legacy_backend = None - CoordinationService._legacy_warning_emitted = False - yield - CoordinationService._legacy_backend = None - CoordinationService._legacy_warning_emitted = False - - def _patch_distributed_coordination(mocker: MockerFixture, backend: object) -> None: mocker.patch( "superset.utils.cache_manager.CacheManager.distributed_coordination", @@ -56,48 +44,39 @@ def test_get_backend_prefers_distributed_coordination( assert CoordinationService.is_backend_defined() is True -def test_get_backend_none_when_nothing_configured( +def test_get_backend_none_when_distributed_coordination_unset( app_context: None, mocker: MockerFixture ) -> None: _patch_distributed_coordination(mocker, None) - mocker.patch.dict( - "flask.current_app.config", {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {}} - ) assert CoordinationService.get_backend() is None assert CoordinationService.is_backend_defined() is False -def test_get_backend_falls_back_to_legacy_gaq_backend_with_warning( +def test_get_backend_ignores_legacy_gaq_config( app_context: None, mocker: MockerFixture ) -> None: + # The coordinator resolves DISTRIBUTED_COORDINATION_CONFIG only; the deprecated + # GAQ backend must never leak into locks/GTF, even with GAQ enabled. _patch_distributed_coordination(mocker, None) mocker.patch("superset.is_feature_enabled", return_value=True) mocker.patch.dict( "flask.current_app.config", {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {"CACHE_TYPE": "RedisCache"}}, ) - legacy_backend = mocker.MagicMock(name="legacy_backend") get_cache_backend = mocker.patch( "superset.async_events.async_query_manager.get_cache_backend", - return_value=legacy_backend, ) - warning = mocker.patch("superset.coordination.base.logger.warning") - assert CoordinationService.get_backend() is legacy_backend - # The legacy backend is memoized and the deprecation warning emitted once. - assert CoordinationService.get_backend() is legacy_backend - get_cache_backend.assert_called_once() - warning.assert_called_once() + assert CoordinationService.get_backend() is None + assert CoordinationService.is_backend_defined() is False + get_cache_backend.assert_not_called() def test_backend_only_ops_raise_when_backend_unavailable( app_context: None, mocker: MockerFixture ) -> None: _patch_distributed_coordination(mocker, None) - mocker.patch.dict( - "flask.current_app.config", {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {}} - ) for op in ( lambda: CoordinationService.publish("channel", "msg"), @@ -111,6 +90,25 @@ def test_backend_only_ops_raise_when_backend_unavailable( op() +def test_ops_use_explicit_backend_without_resolving_coordinator( + app_context: None, mocker: MockerFixture +) -> None: + # An explicit backend (e.g. GAQ's own) is used directly; the shared coordinator + # is never consulted. + get_backend = mocker.patch.object( + CoordinationService, "get_backend", side_effect=AssertionError("resolved") + ) + backend = mocker.MagicMock(name="explicit_backend") + backend.xadd.return_value = "1-0" + + assert ( + CoordinationService.stream_add("stream", {"data": "x"}, backend=backend) + == "1-0" + ) + backend.xadd.assert_called_once_with("stream", {"data": "x"}, "*", None) + get_backend.assert_not_called() + + def test_ops_delegate_to_backend(app_context: None, mocker: MockerFixture) -> None: backend = mocker.MagicMock(name="coordination_backend") backend.publish.return_value = 3 @@ -187,6 +185,22 @@ def test_wait_for_signal_wakes_via_pubsub_and_cleans_up( pubsub.close.assert_called_once() +def test_wait_for_signal_already_satisfied_skips_backend( + app_context: None, mocker: MockerFixture +) -> None: + # An already-satisfied predicate returns straight from the source of truth, + # without resolving or subscribing to a backend — so an already-terminal task + # does not require Redis to be reachable. + backend = mocker.MagicMock(name="backend") + get_backend = mocker.patch.object( + CoordinationService, "get_backend", return_value=backend + ) + + assert CoordinationService.wait_for_signal("ch", lambda: "done") == "done" + get_backend.assert_not_called() + backend.pubsub.assert_not_called() + + def test_listen_invokes_on_signal_then_stops( app_context: None, mocker: MockerFixture ) -> None: From 5e34de66cd52504e4b7d0ba40784eb04f40aa241 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Thu, 20 Aug 2026 17:26:10 -0700 Subject: [PATCH 08/13] refactor(gaq): prefer DISTRIBUTED_COORDINATION_CONFIG, dedicated backend only as fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refines the GAQ backend resolution: AsyncQueryManager now uses the shared coordinator (DISTRIBUTED_COORDINATION_CONFIG) whenever it is configured, and only falls back to its dedicated (deprecated) GLOBAL_ASYNC_QUERIES_CACHE_BACKEND when the coordinator is unset — emitting the one-time deprecation warning only in that fallback case. A deployment that configures DISTRIBUTED_COORDINATION_CONFIG can thus retire the separate GAQ backend instead of maintaining two configs (also covers the Helm chart, which renders a GAQ backend block whenever cache is enabled). Non-breaking: the shipped GLOBAL_ASYNC_QUERIES_CACHE_BACKEND default is unchanged, so enabling GAQ without a coordinator still works. Docs (config.py, UPDATING.md, base.py) updated. --- UPDATING.md | 2 +- superset/async_events/async_query_manager.py | 35 +++++++++++--------- superset/config.py | 24 ++++++++------ superset/coordination/base.py | 7 ++-- 4 files changed, 37 insertions(+), 31 deletions(-) diff --git a/UPDATING.md b/UPDATING.md index 614f31ea98a0..33d0f98942c6 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -61,7 +61,7 @@ the old counter to use the outcome-specific replacements. - [42393](https://github.com/apache/superset/pull/42393): Exported dataset YAML now carries a `uuid` for each metric and column so that custom folder assignments (which reference metrics/columns by UUID) survive an import into another workspace. This affects any export bundle that contains datasets, not just a dataset export: chart, dashboard, database and full-asset exports all embed the same dataset YAML, so a dashboard exported from this release also fails to import into an older one even though no dataset was exported directly. As with `folders` and `currency_code_column`, the affected `datasets/` files fail schema validation (`Unknown field: uuid`) when imported into Superset releases that predate this change; regenerate or hand-edit exports for older targets in mixed-version fleets. - [42300](https://github.com/apache/superset/pull/42300): Timeseries charts (line/area/bar) with a Y-axis bound in effect — either an explicit `yAxisBounds` or one derived from `truncateYAxis` — now clamp out-of-range data points to that bound instead of letting ECharts drop the point (and the line segments around it) entirely. Any existing chart with a configured Y-axis bound and data outside it will look different after upgrading: a gap becomes a point pinned to the boundary. The clamp also rewrites the value ECharts reads for that point's tooltip and data label, so the displayed value is the bound rather than the true observation. - [42087](https://github.com/apache/superset/pull/42087): Stored calculated-column and metric expressions are validated when a query is built, under the same sub-query policy already applied to adhoc expressions. Previously only the dataset update path checked them on save, so expressions written by v1 import, by dataset duplication, or before that check existed were never validated. Since `ALLOW_ADHOC_SUBQUERY` defaults to `False` (see [19242](https://github.com/apache/superset/pull/19242)), a dataset whose stored expression contains a sub-query works before upgrading and afterwards fails at chart render with `Custom SQL fields cannot contain sub-queries.` There is no migration step, and the error does not name the offending dataset column, so audit stored expressions before upgrading: either rewrite them without the sub-query, or set `ALLOW_ADHOC_SUBQUERY = True` to keep the previous behaviour for both stored and adhoc expressions. -- The `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` config key is **deprecated** in favor of the shared coordination backend `DISTRIBUTED_COORDINATION_CONFIG`, which powers a single Redis connection for distributed locks, pub/sub, and the async-events streams (the GAQ firehose). All parameters previously accepted by `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` are supported identically under `DISTRIBUTED_COORDINATION_CONFIG` (they use the same `RedisCache`/`RedisSentinelCache` backend). During the deprecation window the two backends stay scoped: distributed locks and the Global Task Framework use `DISTRIBUTED_COORDINATION_CONFIG` exclusively (falling back to the metadata database when it is unset, unchanged from before), and Global Async Queries keep running on their own `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` (falling back to `DISTRIBUTED_COORDINATION_CONFIG` when that dedicated backend is unset), logging a one-time deprecation warning. This dual-backend arrangement is removed in Superset 8.0, when GAQ moves onto `DISTRIBUTED_COORDINATION_CONFIG`; migrate now by configuring `DISTRIBUTED_COORDINATION_CONFIG`. +- The `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` config key is **deprecated** in favor of the shared coordination backend `DISTRIBUTED_COORDINATION_CONFIG`, which powers a single Redis connection for distributed locks, pub/sub, and the async-events streams (the GAQ firehose). All parameters previously accepted by `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` are supported identically under `DISTRIBUTED_COORDINATION_CONFIG` (they use the same `RedisCache`/`RedisSentinelCache` backend). During the deprecation window the two backends stay scoped: distributed locks and the Global Task Framework use `DISTRIBUTED_COORDINATION_CONFIG` exclusively (falling back to the metadata database when it is unset, unchanged from before), and Global Async Queries use `DISTRIBUTED_COORDINATION_CONFIG` whenever it is set — falling back to the dedicated `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` (with a one-time deprecation warning) only when it is not. Configuring `DISTRIBUTED_COORDINATION_CONFIG` therefore lets a deployment retire the separate `GLOBAL_ASYNC_QUERIES_CACHE_BACKEND` rather than maintain two configs. This dual-backend arrangement is removed in Superset 8.0, when GAQ moves onto `DISTRIBUTED_COORDINATION_CONFIG`; migrate now by configuring it. ### Selenium support removed — Playwright is now required for screenshots diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index 95bcc09ffa38..5b81341607bc 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -159,27 +159,30 @@ def init_app(self, app: Flask) -> None: "is deprecated)." ) - # Global Async Queries keeps its own coordination backend during the - # deprecation window: prefer the dedicated (deprecated) - # GLOBAL_ASYNC_QUERIES_CACHE_BACKEND when configured, otherwise use the - # shared DISTRIBUTED_COORDINATION_CONFIG. This scopes GAQ's stream/pub-sub - # traffic to its own connection and keeps it off the coordinator's backend - # (which powers distributed locks and the Global Task Framework). In 8.0 the - # dedicated backend is removed and GAQ moves onto the coordinator's connection. - if app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get("CACHE_TYPE"): + # Global Async Queries share the coordinator's backend whenever + # DISTRIBUTED_COORDINATION_CONFIG is configured, so an 8.0-leaning deployment + # keeps a single coordination config rather than maintaining a separate one. + # Only when no coordinator is configured does GAQ fall back to its dedicated + # (deprecated) GLOBAL_ASYNC_QUERIES_CACHE_BACKEND, emitting a one-time + # deprecation warning. That dedicated backend is removed in Superset 8.0, when + # DISTRIBUTED_COORDINATION_CONFIG becomes the only option. Either way GAQ passes + # its resolved backend explicitly to CoordinationService's primitives, keeping + # its stream/pub-sub traffic scoped to the connection it resolved here. + if (coordinator := CoordinationService.get_backend()) is not None: + self._gaq_backend = coordinator + else: if not AsyncQueryManager._legacy_backend_warning_emitted: logger.warning( - "Global Async Queries is running on its own " - "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND. This dedicated backend is " - "deprecated and will be removed in Superset 8.0, when GAQ will " - "use DISTRIBUTED_COORDINATION_CONFIG like the rest of Superset's " - "coordination (distributed locks, task framework, pub/sub). " - "Configure DISTRIBUTED_COORDINATION_CONFIG to consolidate now." + "Global Async Queries is running on the deprecated " + "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND because " + "DISTRIBUTED_COORDINATION_CONFIG is not configured. Configure " + "DISTRIBUTED_COORDINATION_CONFIG to consolidate coordination " + "(distributed locks, task framework, pub/sub, and the GAQ event " + "streams) onto a single connection; the dedicated backend is " + "removed in Superset 8.0." ) AsyncQueryManager._legacy_backend_warning_emitted = True self._gaq_backend = get_cache_backend(app.config) - else: - self._gaq_backend = CoordinationService.get_backend() if len(app.config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"]) < 32: raise AsyncQueryTokenException( diff --git a/superset/config.py b/superset/config.py index b6b3d814e56c..eea684e1793d 100644 --- a/superset/config.py +++ b/superset/config.py @@ -2943,13 +2943,14 @@ def EMAIL_HEADER_MUTATOR( # pylint: disable=invalid-name,unused-argument # noq # - Set 'CACHE_TYPE' to 'RedisSentinelCache' for RedisSentinelCacheBackend. # # DEPRECATED: this dedicated backend is retained only so Global Async Queries can -# keep running on its own coordination connection during the deprecation window. -# When configured it is used by GAQ *only* (its event streams and cancel registry); -# it never powers distributed locks or the Global Task Framework, which use -# DISTRIBUTED_COORDINATION_CONFIG exclusively. GAQ prefers this dedicated backend -# when set and otherwise falls back to DISTRIBUTED_COORDINATION_CONFIG. This -# dual-backend arrangement is deprecated and removed in Superset 8.0, when GAQ -# moves onto DISTRIBUTED_COORDINATION_CONFIG like the rest of Superset's +# keep running on their own coordination connection when DISTRIBUTED_COORDINATION_CONFIG +# is not configured. When configured it is used by GAQ *only* (its event streams and +# cancel registry); it never powers distributed locks or the Global Task Framework, +# which use DISTRIBUTED_COORDINATION_CONFIG exclusively. GAQ uses +# DISTRIBUTED_COORDINATION_CONFIG whenever it is set and only falls back to this +# dedicated backend when it is not, so a consolidated deployment need not maintain a +# second config. This dual-backend arrangement is deprecated and removed in Superset +# 8.0, when GAQ moves onto DISTRIBUTED_COORDINATION_CONFIG like the rest of Superset's # coordination. All parameters here are supported identically under # DISTRIBUTED_COORDINATION_CONFIG (both use the same # RedisCache/RedisSentinelCache backend). @@ -3246,10 +3247,11 @@ class ExtraAccessQueryFilters(TypedDict, total=False): # (``superset.coordination.base.CoordinationService``) exposing standardized interfaces # for distributed locks, pub/sub, and streams under a single connection. It is the # single source of truth for the coordinator's consumers (distributed locks, the -# Global Task Framework, and future stream/pub-sub users). Global Async Queries keep -# their own dedicated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` for now (falling back to -# this when that is unset); that dual-backend arrangement is deprecated and, in -# Superset 8.0, GAQ moves onto this connection and the dedicated backend is removed. +# Global Task Framework, and future stream/pub-sub users). Global Async Queries use +# this connection whenever it is set, falling back to their dedicated +# ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` only when it is unset; that dual-backend +# arrangement is deprecated and, in Superset 8.0, GAQ moves onto this connection and +# the dedicated backend is removed. # # All parameters previously supported by GLOBAL_ASYNC_QUERIES_CACHE_BACKEND are # supported here (both go through the same RedisCacheBackend/RedisSentinelCacheBackend diff --git a/superset/coordination/base.py b/superset/coordination/base.py index 2f8961b76692..d6b783f1b013 100644 --- a/superset/coordination/base.py +++ b/superset/coordination/base.py @@ -78,9 +78,10 @@ def get_backend(cls) -> "CoordinationBackend | None": the single source of truth for the coordinator's consumers (distributed locks, the Global Task Framework, and future stream/pub-sub users); it does *not* consult the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND``. Global - Async Queries owns its own separate backend during the deprecation window — - see :class:`~superset.async_events.async_query_manager.AsyncQueryManager` — - and passes it explicitly to the primitives below via ``backend``. + Async Queries resolve their own backend (this coordinator when configured, + else the deprecated dedicated backend — see + :class:`~superset.async_events.async_query_manager.AsyncQueryManager`) and pass + it explicitly to the primitives below via ``backend``. """ from superset.extensions import cache_manager From 417087c0bffc79e163cd3a6c9050661851e054ba Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Fri, 21 Aug 2026 08:12:11 -0700 Subject: [PATCH 09/13] fix(coordination): close pub/sub connection even if unsubscribe fails close_pubsub now attempts unsubscribe() and close() independently, so a failing unsubscribe no longer skips close() and leak the underlying Redis/pub-sub connection (the helper runs from every listener teardown and wait_for_signal finally block). Also corrects the CoordinationService docstring: the await/notify layer wakes on a published message when a backend is defined and polls only when none is defined; it does not poll as an outage fallback once a backend is configured. --- superset/coordination/base.py | 7 ++++--- superset/coordination/utils.py | 9 ++++++++- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/superset/coordination/base.py b/superset/coordination/base.py index d6b783f1b013..af94b036b47b 100644 --- a/superset/coordination/base.py +++ b/superset/coordination/base.py @@ -57,9 +57,10 @@ class CoordinationService: - **Higher-level await/notify** — ``wait_for_signal`` (blocking) and ``listen_for_signal`` (background). These combine a pub/sub channel with a caller-supplied predicate: - when a backend is defined they wake promptly on a published message, and either - way they fall back to polling the predicate. This keeps the pub/sub-vs-poll - boilerplate in one place; callers just supply a channel and a check. + when a backend is defined they wake promptly on a published message and + re-check the predicate each tick; without a backend they poll the predicate. + This keeps the pub/sub-vs-poll boilerplate in one place; callers just supply a + channel and a check. All methods are class-level: the service is app-global and resolves its backend from the shared coordination connection on each call. diff --git a/superset/coordination/utils.py b/superset/coordination/utils.py index 7a567c6b7d31..153702cd72c8 100644 --- a/superset/coordination/utils.py +++ b/superset/coordination/utils.py @@ -25,9 +25,16 @@ def close_pubsub(pubsub: Any) -> None: - """Best-effort unsubscribe + close of a pub/sub subscription.""" + """Best-effort unsubscribe + close of a pub/sub subscription. + + ``unsubscribe`` and ``close`` are attempted independently so a failure of the + former still releases the underlying connection. + """ try: pubsub.unsubscribe() + except Exception as ex: # pylint: disable=broad-except + logger.debug("Error unsubscribing pub/sub subscription: %s", ex) + try: pubsub.close() except Exception as ex: # pylint: disable=broad-except logger.debug("Error closing pub/sub subscription: %s", ex) From 3958e2463611813d37de17ca1d8efcae96752c27 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Fri, 21 Aug 2026 08:51:07 -0700 Subject: [PATCH 10/13] test(coordination): update wait_for_signal tests for pre-subscribe check The wait_for_signal fast-path fix evaluates check() once before opening the pub/sub subscription, so the mocked tests that assumed the first check happens after subscribe needed one extra predicate value: test_wait_for_signal_wakes_via_pubsub_and_cleans_up (add a None so a pub/sub nudge still occurs) and TaskManager's test_pubsub_success_subscribes_and_cleans_up (add a pending read so we still subscribe). --- tests/unit_tests/coordination/test_service.py | 6 ++++-- tests/unit_tests/tasks/test_manager.py | 4 +++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py index 49250492c11a..c665ae37a9b0 100644 --- a/tests/unit_tests/coordination/test_service.py +++ b/tests/unit_tests/coordination/test_service.py @@ -172,14 +172,16 @@ def test_wait_for_signal_wakes_via_pubsub_and_cleans_up( pubsub = mocker.MagicMock(name="pubsub") backend.pubsub.return_value = pubsub mocker.patch.object(CoordinationService, "get_backend", return_value=backend) - results = iter([None, "done"]) + # First check (pre-subscribe fast path) is None, so we subscribe; the next None + # forces one pub/sub nudge, then "done" satisfies the wait. + results = iter([None, None, "done"]) result = CoordinationService.wait_for_signal( "ch", lambda: next(results), timeout=5.0 ) assert result == "done" pubsub.subscribe.assert_called_once_with("ch") - # One wake-up nudge between the first (None) and second (done) check. + # One wake-up nudge between the post-subscribe (None) and final (done) check. pubsub.get_message.assert_called_once() pubsub.unsubscribe.assert_called_once() pubsub.close.assert_called_once() diff --git a/tests/unit_tests/tasks/test_manager.py b/tests/unit_tests/tasks/test_manager.py index 96afe46fd6d2..5a309f7d8fa0 100644 --- a/tests/unit_tests/tasks/test_manager.py +++ b/tests/unit_tests/tasks/test_manager.py @@ -247,7 +247,9 @@ def test_pubsub_success_subscribes_and_cleans_up(self, mock_dao, mock_get_backen pending.status = "pending" complete = MagicMock() complete.status = "success" - mock_dao.find_one_or_none.side_effect = [pending, complete] + # find_one_or_none is called for: the existence check, the pre-subscribe fast + # path (still pending → we subscribe), then the post-subscribe check (success). + mock_dao.find_one_or_none.side_effect = [pending, pending, complete] backend = MagicMock() pubsub = MagicMock() From ae62d1db9c4bf64f8daeab0c12e40e5f0728f26b Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Fri, 21 Aug 2026 09:02:55 -0700 Subject: [PATCH 11/13] docs(gaq): correct cancel-registry comment to match actual behavior The write is skipped only when no coordination backend is configured; when one is configured a write failure propagates and fails job submission (it is not swallowed). The prior 'best-effort' wording overstated the guarantee. --- superset/async_events/async_query_manager.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index 5b81341607bc..a99643efa35c 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -342,8 +342,10 @@ def _register_cancellable_job( owner without trusting the client-supplied id. Expires with the JWT so it never outlives the job it guards. """ - # Best-effort: the cancel registry is an optimization. Skip when no - # coordination backend is configured rather than failing job submission. + # The cancel registry is an optimization: skip it when GAQ has no + # coordination backend configured. When a backend is configured, a write + # failure is not swallowed here — it surfaces to the caller and fails job + # submission, matching the behavior of the surrounding stream writes. if self._gaq_backend is None: return CoordinationService.set_value( From def9996053f39ce5efd60006b677e610232ad406 Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Fri, 21 Aug 2026 09:04:34 -0700 Subject: [PATCH 12/13] fix(gaq): gate is_job_cancelled on its own backend; clarify service docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit is_job_cancelled was the one GAQ method missing the self._gaq_backend is None guard, so with no dedicated backend it would fall through _require_backend() to the shared coordinator — the cross-backend resolution GAQ is meant to avoid. Restores the guard (matching the pre-refactor behavior of returning False). Also corrects the CoordinationService class docstring to note the primitives accept an explicit backend rather than always using the shared connection. --- superset/async_events/async_query_manager.py | 2 ++ superset/coordination/base.py | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py index a99643efa35c..bdf94c4fb452 100644 --- a/superset/async_events/async_query_manager.py +++ b/superset/async_events/async_query_manager.py @@ -463,6 +463,8 @@ def is_job_cancelled(self, job_id: str) -> bool: swallowed and treated as "not cancelled" — a Redis blip must never mask the original error (e.g. a genuine timeout) with a connection error. """ + if self._gaq_backend is None: + return False try: raw = CoordinationService.get_value( self._job_registry_key(job_id), backend=self._gaq_backend diff --git a/superset/coordination/base.py b/superset/coordination/base.py index af94b036b47b..996a1f494e2f 100644 --- a/superset/coordination/base.py +++ b/superset/coordination/base.py @@ -62,8 +62,10 @@ class CoordinationService: This keeps the pub/sub-vs-poll boilerplate in one place; callers just supply a channel and a check. - All methods are class-level: the service is app-global and resolves its backend - from the shared coordination connection on each call. + All methods are class-level: the service is app-global. Calls resolve the shared + coordination backend from ``DISTRIBUTED_COORDINATION_CONFIG`` on each call, except + the raw primitives, which accept an optional ``backend`` so a caller with its own + connection (Global Async Queries, during the deprecation window) can run against it. Distributed locking is *not* exposed here: it has its own user-facing interface (:class:`~superset.distributed_lock.DistributedLock`) that uses this service's From 2d81d1cb5e8fa9fa16be9330be1d2a2dbcc5a23b Mon Sep 17 00:00:00 2001 From: Ville Brofeldt Date: Fri, 21 Aug 2026 10:29:13 -0700 Subject: [PATCH 13/13] test(async-events): init_app before login in run_test_with_cache_backend The GAQ backend must be resolved under the get_cache_backend patch, but init_app also re-registers the after_request handler and must run before login(): with login() first, its request tripped Flask's 'setup method after first request' guard before init_app re-registered the handler, failing all 6 TestAsyncEventApi cases. Resolve the mock under the patch during init_app (the backend is captured there), then log in and run. --- tests/integration_tests/async_events/api_tests.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/integration_tests/async_events/api_tests.py b/tests/integration_tests/async_events/api_tests.py index c711453c550c..89e008940bd6 100644 --- a/tests/integration_tests/async_events/api_tests.py +++ b/tests/integration_tests/async_events/api_tests.py @@ -48,16 +48,18 @@ def run_test_with_cache_backend(self, cache_backend_cls: Type[Any], test_func): app._got_first_request = False # GAQ resolves its own backend from get_cache_backend during init_app, so - # inject the mock there and initialize within the patch. + # inject the mock there. init_app must run before login(): it re-registers the + # after_request handler, and login()'s request would otherwise trip Flask's + # "setup method after first request" guard before init_app gets to re-register. mock_cache = mock.Mock(spec=cache_backend_cls) - - self.login(ADMIN_USERNAME) with mock.patch( "superset.async_events.async_query_manager.get_cache_backend", return_value=mock_cache, ): async_query_manager_factory.init_app(app) - test_func(mock_cache) + + self.login(ADMIN_USERNAME) + test_func(mock_cache) def _test_events_logic(self, mock_cache): with mock.patch.object(mock_cache, "xrange") as mock_xrange: