(superset) 01/01: fix(tasks): let deployments propagate context into the abort-listener thread

[el...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

eschutho pushed a commit to branch fix-task-abort-listener-context
in repository https://gitbox.apache.org/repos/asf/superset.git

commit cf2fd0f0481a52abf83818aa62cad1a128beacc8
Author: Elizabeth Thompson <es...@gmail.com>
AuthorDate: Mon Oct 19 18:50:09 2026 +0000

    fix(tasks): let deployments propagate context into the abort-listener thread
    
    The task abort listener polls from a background thread. Any deployment that
    keeps per-request state outside the app context (e.g. a multi-tenant
    workspace binding on flask.g) has none of it there, so every poll of
    _check_abort_status fails and abort signals are never delivered.
    
    Add an opt-in BACKGROUND_THREAD_CONTEXT_FACTORY config hook. It is resolved
    on the calling thread and the context manager it yields wraps the listener
    thread's lifetime. Default is a no-op, so behaviour is unchanged.
    
    Co-Authored-By: Claude Sonnet 5.5 <no...@anthropic.com>
---
 superset/config.py                            | 23 +++++++++++++-
 superset/coordination/base.py                 | 27 ++++++++++++++--
 tests/unit_tests/coordination/test_service.py | 46 +++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 4 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index ceb47817f78..4ea0d090167 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -37,7 +37,16 @@ from contextlib import contextmanager
 from datetime import timedelta
 from email.mime.multipart import MIMEMultipart
 from importlib.resources import files
-from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, TypedDict
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Iterator,
+    Literal,
+    Optional,
+    TYPE_CHECKING,
+    TypedDict,
+)
 
 import click
 from celery.schedules import crontab
@@ -3653,6 +3662,18 @@ DISTRIBUTED_COORDINATION_CONFIG: CacheConfig | None = None
 # indefinitely. Defaults to 24 hours.
 DISTRIBUTED_COORDINATION_SIGNAL_TTL = int(timedelta(hours=24).total_seconds())
 
+# Optional hook for propagating caller-thread context (e.g. a multi-tenant
+# deployment's per-workspace state) into background threads Superset spawns, such as
+# the task abort listener. Superset calls the factory on the *calling* thread, where
+# that context is live; the factory returns a zero-argument callable that the
+# background thread invokes to get a context manager, which wraps the thread's whole
+# lifetime. The context manager is responsible for pushing its own app context
+# (never share the caller's: SQLAlchemy sessions are scoped to it). Defaults to
+# None, which leaves background threads to push a bare app context as before.
+BACKGROUND_THREAD_CONTEXT_FACTORY: (
+    Callable[[], Callable[[], ContextManager[Any]]] | None
+) = None
+
 # Maximum time to wait for a semantic-cache descriptor lease.
 SEMANTIC_CACHE_COORDINATION_WAIT_SECONDS: float = 1.0
 
diff --git a/superset/coordination/base.py b/superset/coordination/base.py
index 85bbba6b724..1bb66c24ac5 100644
--- a/superset/coordination/base.py
+++ b/superset/coordination/base.py
@@ -21,10 +21,11 @@ See :mod:`superset.coordination` for the package overview.
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import threading
 import time
-from typing import Any, Callable, TYPE_CHECKING, TypeVar, Union
+from typing import Any, Callable, ContextManager, TYPE_CHECKING, TypeVar, Union
 
 from flask import current_app
 from redis.exceptions import RedisError, TimeoutError as RedisTimeoutError
@@ -497,9 +498,16 @@ class CoordinationService:
         :param name: optional thread name suffix for logging.
         """
         stop_event = threading.Event()
+        enter_context = cls._capture_background_context()
+
+        def run() -> None:
+            with enter_context():
+                cls._run_listen_loop(
+                    channel, check, on_signal, stop_event, poll_interval
+                )
+
         thread = threading.Thread(
-            target=cls._run_listen_loop,
-            args=(channel, check, on_signal, stop_event, poll_interval),
+            target=run,
             daemon=True,
             name=f"coord-listen-{name or channel}",
         )
@@ -525,6 +533,19 @@ class CoordinationService:
 
         return SignalListener(thread, stop_event, wake=wake)
 
+    @staticmethod
+    def _capture_background_context() -> Callable[[], ContextManager[Any]]:
+        """Snapshot caller-thread context for a background thread to re-enter.
+
+        Resolves ``BACKGROUND_THREAD_CONTEXT_FACTORY`` on the calling thread (where
+        the request/task context is live) and returns the callable the background
+        thread uses to enter an equivalent context. Without a configured factory
+        the returned context manager is a no-op, leaving the thread to push a bare
+        app context per call as before.
+        """
+        factory = current_app.config.get("BACKGROUND_THREAD_CONTEXT_FACTORY")
+        return factory() if factory else contextlib.nullcontext
+
     @classmethod
     def _run_listen_loop(
         cls,
diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py
index 2f8f1e43022..4421b65601d 100644
--- a/tests/unit_tests/coordination/test_service.py
+++ b/tests/unit_tests/coordination/test_service.py
@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import threading
+from typing import Callable
 
 import pytest
 from pytest_mock import MockerFixture
@@ -456,3 +457,48 @@ def test_listen_stop_wakes_blocked_backend_read(
         call.args[:2] == ("ch", {"m": "__wake__"})
         for call in backend.xadd.call_args_list
     )
+
+
+def test_listen_runs_inside_configured_background_context(
+    app_context: None, mocker: MockerFixture
+) -> None:
+    """The factory runs on the caller thread; its context wraps the listener thread."""
+    from flask import current_app
+
+    mocker.patch.object(CoordinationService, "get_backend", return_value=None)
+    caller = threading.current_thread()
+    events: list[str] = []
+    factory_threads: list[threading.Thread] = []
+
+    class _Ctx:
+        def __enter__(self) -> None:
+            events.append("enter")
+
+        def __exit__(self, *exc: object) -> None:
+            events.append("exit")
+
+    def factory() -> Callable[[], _Ctx]:
+        factory_threads.append(threading.current_thread())
+        return _Ctx
+
+    current_app.config["BACKGROUND_THREAD_CONTEXT_FACTORY"] = factory
+    try:
+        checked_inside: list[bool] = []
+
+        def check() -> bool:
+            checked_inside.append(events == ["enter"])
+            return True
+
+        fired = threading.Event()
+        listener = CoordinationService.listen_for_signal(
+            "ch", check=check, on_signal=fired.set, poll_interval=0.01
+        )
+        assert fired.wait(timeout=2.0) is True
+        listener.stop()
+        listener._thread.join(timeout=2.0)
+    finally:
+        current_app.config.pop("BACKGROUND_THREAD_CONTEXT_FACTORY", None)
+
+    assert factory_threads == [caller]
+    assert checked_inside == [True]
+    assert events == ["enter", "exit"]