You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by vi...@apache.org on 2026/08/23 17:57:25 UTC
(superset) branch villebro/coordination-kv-callable-key created (now c6cc291ade1)
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a change to branch villebro/coordination-kv-callable-key
in repository https://gitbox.apache.org/repos/asf/superset.git
at c6cc291ade1 feat(coordination): accept a callable key generator in the KV API
This branch includes the following new commits:
new c6cc291ade1 feat(coordination): accept a callable key generator in the KV API
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: feat(coordination): accept a callable key generator in the KV API
Posted by vi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a commit to branch villebro/coordination-kv-callable-key
in repository https://gitbox.apache.org/repos/asf/superset.git
commit c6cc291ade1b3bdd93659b8e2878d97ec1b09ed4
Author: Ville Brofeldt <vi...@gmail.com>
AuthorDate: Sun Aug 23 10:56:51 2026 -0700
feat(coordination): accept a callable key generator in the KV API
CoordinationService.get_value / set_value / delete_value took a plain key
string; the caller always built the key eagerly. For parity with Superset's
other cache-key helpers (e.g. memoized_func(key=...)), the KV ops now also
accept a () -> str key *generator*, resolved at call time via a shared
_resolve_key helper (a non-string result raises TypeError). A literal string
still passes through unchanged, so this is backward-compatible and additive.
Note this is scoped to the general-purpose coordinator KV only — the GTF
task_key stays a fixed literal (it is the dedup identity resolved synchronously
at submit time, feeding the dedup lock + dedup_key).
---
superset/coordination/base.py | 47 +++++++++++++++++++++++----
tests/unit_tests/coordination/test_service.py | 39 ++++++++++++++++++++++
2 files changed, 79 insertions(+), 7 deletions(-)
diff --git a/superset/coordination/base.py b/superset/coordination/base.py
index 7a1c2361e70..e05f63ea091 100644
--- a/superset/coordination/base.py
+++ b/superset/coordination/base.py
@@ -24,7 +24,7 @@ from __future__ import annotations
import logging
import threading
import time
-from typing import Any, Callable, TYPE_CHECKING, TypeVar
+from typing import Any, Callable, TYPE_CHECKING, TypeVar, Union
from flask import current_app
from redis.exceptions import TimeoutError as RedisTimeoutError
@@ -39,6 +39,12 @@ logger = logging.getLogger(__name__)
T = TypeVar("T")
+# A coordinator KV key may be a literal string or a zero-argument generator
+# resolved at call time, for parity with Superset's other cache-key helpers
+# (e.g. ``memoized_func(key=...)``). Callers that build a key lazily (or want a
+# single shared key-derivation callable) can pass the callable directly.
+KeyLike = Union[str, Callable[[], str]]
+
# Reliable signalling uses Redis Streams: entries are persisted and replayable, so a
# waiter still receives a signal if it subscribes slightly late, reconnects, or
# survives a failover. A signal is a single stream entry; the caller-supplied
@@ -160,19 +166,40 @@ class CoordinationService:
# -- Key/Value -----------------------------------------------------------
+ @staticmethod
+ def _resolve_key(key: KeyLike) -> str:
+ """Resolve a KV key that may be a literal string or a ``() -> str`` generator.
+
+ Lets callers pass a key generator (resolved here, at call time) for parity
+ with Superset's other cache-key helpers, while a plain string passes through
+ unchanged.
+
+ :raises TypeError: if a generator returns a non-string.
+ """
+ resolved = key() if callable(key) else key
+ if not isinstance(resolved, str):
+ raise TypeError(
+ f"Coordinator KV key must resolve to str, got {type(resolved).__name__}"
+ )
+ return resolved
+
@classmethod
- def get_value(cls, key: str, backend: "CoordinationBackend | None" = None) -> Any:
+ def get_value(
+ cls, key: KeyLike, backend: "CoordinationBackend | None" = None
+ ) -> Any:
"""Return the raw (bytes) value at ``key``, or ``None`` if absent.
+ :param key: a literal key string, or a ``() -> str`` generator (see
+ :meth:`_resolve_key`).
:param backend: optional explicit backend (see :meth:`_require_backend`).
:raises CoordinationBackendUnavailableError: if no backend is available.
"""
- return cls._require_backend(backend).get(key)
+ return cls._require_backend(backend).get(cls._resolve_key(key))
@classmethod
def set_value(
cls,
- key: str,
+ key: KeyLike,
value: Any,
ttl: int | None = None,
if_absent: bool = False,
@@ -181,6 +208,8 @@ class CoordinationService:
) -> bool | None:
"""Store ``value`` at ``key``.
+ :param key: a literal key string, or a ``() -> str`` generator (see
+ :meth:`_resolve_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.
@@ -190,19 +219,23 @@ class CoordinationService:
:raises CoordinationBackendUnavailableError: if no backend is available.
"""
return cls._require_backend(backend).set(
- key, value, ex=ttl, nx=if_absent, xx=if_present
+ cls._resolve_key(key), value, ex=ttl, nx=if_absent, xx=if_present
)
@classmethod
def delete_value(
- cls, *keys: str, backend: "CoordinationBackend | None" = None
+ cls, *keys: KeyLike, backend: "CoordinationBackend | None" = None
) -> int:
"""Delete one or more keys; returns the number deleted.
+ :param keys: literal key strings and/or ``() -> str`` generators (see
+ :meth:`_resolve_key`).
:param backend: optional explicit backend (see :meth:`_require_backend`).
:raises CoordinationBackendUnavailableError: if no backend is available.
"""
- return cls._require_backend(backend).delete(*keys)
+ return cls._require_backend(backend).delete(
+ *(cls._resolve_key(key) for key in keys)
+ )
# -- Streams -------------------------------------------------------------
diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py
index cc9419a2b84..98d1ba4e989 100644
--- a/tests/unit_tests/coordination/test_service.py
+++ b/tests/unit_tests/coordination/test_service.py
@@ -116,6 +116,45 @@ def test_ops_delegate_to_backend(app_context: None, mocker: MockerFixture) -> No
backend.xrange.assert_called_once_with("stream", "-", "+", 10)
+def test_kv_ops_accept_a_callable_key(app_context: None, mocker: MockerFixture) -> None:
+ """A key may be a ``() -> str`` generator, resolved at call time."""
+ backend = mocker.MagicMock(name="coordination_backend")
+ backend.get.return_value = b"v"
+ backend.set.return_value = True
+ backend.delete.return_value = 2
+ _patch_distributed_coordination(mocker, backend)
+
+ calls = {"n": 0}
+
+ def key_gen() -> str:
+ calls["n"] += 1
+ return "generated-key"
+
+ CoordinationService.get_value(key_gen)
+ CoordinationService.set_value(key_gen, "v")
+ # A mix of literal and callable keys resolves each independently.
+ CoordinationService.delete_value("literal-key", key_gen)
+
+ backend.get.assert_called_once_with("generated-key")
+ backend.set.assert_called_once_with(
+ "generated-key", "v", ex=None, nx=False, xx=False
+ )
+ backend.delete.assert_called_once_with("literal-key", "generated-key")
+ # The generator was invoked once per op (resolved lazily, not cached).
+ assert calls["n"] == 3
+
+
+def test_kv_key_generator_returning_non_string_raises(
+ app_context: None, mocker: MockerFixture
+) -> None:
+ backend = mocker.MagicMock(name="coordination_backend")
+ _patch_distributed_coordination(mocker, backend)
+
+ with pytest.raises(TypeError):
+ CoordinationService.get_value(lambda: 123) # type: ignore[arg-type, return-value]
+ backend.get.assert_not_called()
+
+
# -- wait_for_signal / listen --------------------------------------------------