From 040343c7edd9a6072f28d1c448163d12aa2f2c97 Mon Sep 17 00:00:00 2001 From: Wei Chen <182598863+cw-jlu@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:30:12 +0800 Subject: [PATCH 1/3] Read NATS schedules with multi-token identifiers --- taskiq_nats/schedule_source.py | 17 +++-- tests/test_schedule_ids.py | 131 +++++++++++++++++++++++++++++++++ 2 files changed, 142 insertions(+), 6 deletions(-) create mode 100644 tests/test_schedule_ids.py diff --git a/taskiq_nats/schedule_source.py b/taskiq_nats/schedule_source.py index 4855842..5e26a00 100644 --- a/taskiq_nats/schedule_source.py +++ b/taskiq_nats/schedule_source.py @@ -96,15 +96,20 @@ async def get_schedules(self) -> list[ScheduledTask]: :return: list of schedules. """ try: - schedules = await self.kv.history(f"{self.prefix}.*") + schedules = await self.kv.history(f"{self.prefix}.>") except NoKeysError: return [] - return [ - ScheduledTask.model_validate(self.serializer.loadb(schedule.value)) - for schedule in schedules - if schedule and schedule.value - ] + result = [] + for schedule in schedules: + if schedule and schedule.value: + task = ScheduledTask.model_validate( + self.serializer.loadb(schedule.value), + ) + # A nested prefix can match the wildcard without belonging to us. + if schedule.key == f"{self.prefix}.{task.schedule_id}": + result.append(task) + return result async def post_send(self, task: ScheduledTask) -> None: """Delete a task after it's completed.""" diff --git a/tests/test_schedule_ids.py b/tests/test_schedule_ids.py new file mode 100644 index 0000000..fb5c267 --- /dev/null +++ b/tests/test_schedule_ids.py @@ -0,0 +1,131 @@ +import datetime as dt +import pickle +import uuid +from collections.abc import AsyncGenerator + +import pytest +from taskiq import ScheduledTask +from taskiq.cli.scheduler.run import get_schedules + +from taskiq_nats import NATSKeyValueScheduleSource + + +@pytest.fixture +async def schedule_source( + nats_urls: list[str], +) -> AsyncGenerator[NATSKeyValueScheduleSource, None]: + source = NATSKeyValueScheduleSource( + servers=nats_urls, + bucket_name=f"schedules_{uuid.uuid4().hex}", + prefix="schedules", + ) + await source.startup() + try: + yield source + finally: + try: + await source.nats_jetstream.delete_key_value(source.bucket_name) + finally: + await source.shutdown() + + +@pytest.mark.parametrize("prefix", ["schedules", "team.schedules"]) +@pytest.mark.parametrize("kind", ["cron", "time"]) +async def test_dotted_schedule_ids_round_trip( + nats_urls: list[str], + schedule_source: NATSKeyValueScheduleSource, + prefix: str, + kind: str, +) -> None: + source = NATSKeyValueScheduleSource( + servers=nats_urls, + bucket_name=schedule_source.bucket_name, + prefix=prefix, + ) + await source.startup() + try: + schedules = [ + ScheduledTask( + schedule_id=name, + task_name="report", + labels={}, + args=[], + kwargs={}, + cron="* * * * *" if kind == "cron" else None, + time=( + dt.datetime(2030, 1, 1, tzinfo=dt.timezone.utc) + if kind == "time" + else None + ), + ) + for name in ["plain", "daily.report", "team.daily.report"] + ] + for schedule in schedules: + await source.add_schedule(schedule) + + assert sorted(await source.kv.keys()) == sorted( + f"{prefix}.{task.schedule_id}" for task in schedules + ) + assert {task.schedule_id: task for task in await source.get_schedules()} == { + task.schedule_id: task for task in schedules + } + assert {task.schedule_id for task in await get_schedules(source)} == { + task.schedule_id for task in schedules + } + + await source.delete_schedule("daily.report") + assert {task.schedule_id for task in await source.get_schedules()} == { + "plain", + "team.daily.report", + } + restored = schedules[1].model_copy(update={"args": ["updated"]}) + await source.add_schedule(restored) + snapshot = {task.schedule_id: task for task in await source.get_schedules()} + assert snapshot[restored.schedule_id] == restored + + await source.post_send(restored) + assert ( + restored.schedule_id + in {task.schedule_id for task in await get_schedules(source)} + ) == (kind == "cron") + finally: + await source.shutdown() + + +async def test_nested_prefix_keeps_its_own_schedules( + nats_urls: list[str], + schedule_source: NATSKeyValueScheduleSource, +) -> None: + nested = NATSKeyValueScheduleSource( + servers=nats_urls, + bucket_name=schedule_source.bucket_name, + prefix=f"{schedule_source.prefix}.nested", + ) + await nested.startup() + try: + root_task = ScheduledTask( + schedule_id="root_task", + task_name="report", + labels={}, + args=[], + kwargs={}, + cron="* * * * *", + ) + nested_task = root_task.model_copy(update={"schedule_id": "nested_task"}) + await schedule_source.add_schedule(root_task) + await nested.add_schedule(nested_task) + assert await schedule_source.get_schedules() == [root_task] + assert await nested.get_schedules() == [nested_task] + finally: + await nested.shutdown() + + +async def test_invalid_payload_still_raises_and_recovers( + schedule_source: NATSKeyValueScheduleSource, +) -> None: + key = f"{schedule_source.prefix}.broken" + await schedule_source.kv.put(key, b"not-pickle") + with pytest.raises(pickle.UnpicklingError): + await schedule_source.get_schedules() + await schedule_source.kv.delete(key) + assert await schedule_source.get_schedules() == [] From da0f9c5b7da4241c888271f4e667f859c19014e4 Mon Sep 17 00:00:00 2001 From: Wei Chen <182598863+cw-jlu@users.noreply.github.com> Date: Fri, 13 Nov 2026 20:59:17 +0800 Subject: [PATCH 2/3] Keep nested-prefix payloads from breaking the snapshot Reading prefix.> also returns keys written by sources with a nested prefix, which may use another serializer. Skip and log an undecodable multi-token key instead of failing the whole snapshot; such keys were never read before multi-token IDs were supported. Undecodable single-token keys can only belong to this source and still raise. Co-Authored-By: Claude Opus 5.5 --- taskiq_nats/schedule_source.py | 19 ++++++++++-- tests/test_schedule_ids.py | 56 ++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 3 deletions(-) diff --git a/taskiq_nats/schedule_source.py b/taskiq_nats/schedule_source.py index 5e26a00..d7e11de 100644 --- a/taskiq_nats/schedule_source.py +++ b/taskiq_nats/schedule_source.py @@ -103,9 +103,22 @@ async def get_schedules(self) -> list[ScheduledTask]: result = [] for schedule in schedules: if schedule and schedule.value: - task = ScheduledTask.model_validate( - self.serializer.loadb(schedule.value), - ) + try: + task = ScheduledTask.model_validate( + self.serializer.loadb(schedule.value), + ) + except Exception: + # A multi-token key may belong to a source with a nested + # prefix and a different serializer, so it must not break + # this snapshot. Single-token keys can only be ours. + if "." not in schedule.key[len(self.prefix) + 1 :]: + raise + log.warning( + "Skipping undecodable schedule %r: it may belong to " + "a source with a nested prefix.", + schedule.key, + ) + continue # A nested prefix can match the wildcard without belonging to us. if schedule.key == f"{self.prefix}.{task.schedule_id}": result.append(task) diff --git a/tests/test_schedule_ids.py b/tests/test_schedule_ids.py index fb5c267..85e05ad 100644 --- a/tests/test_schedule_ids.py +++ b/tests/test_schedule_ids.py @@ -1,4 +1,5 @@ import datetime as dt +import logging import pickle import uuid from collections.abc import AsyncGenerator @@ -6,6 +7,7 @@ import pytest from taskiq import ScheduledTask from taskiq.cli.scheduler.run import get_schedules +from taskiq.serializers import JSONSerializer from taskiq_nats import NATSKeyValueScheduleSource @@ -120,6 +122,38 @@ async def test_nested_prefix_keeps_its_own_schedules( await nested.shutdown() +async def test_nested_prefix_with_another_serializer( + nats_urls: list[str], + schedule_source: NATSKeyValueScheduleSource, + caplog: pytest.LogCaptureFixture, +) -> None: + nested = NATSKeyValueScheduleSource( + servers=nats_urls, + bucket_name=schedule_source.bucket_name, + prefix=f"{schedule_source.prefix}.nested", + serializer=JSONSerializer(), + ) + await nested.startup() + try: + root_task = ScheduledTask( + schedule_id="root_task", + task_name="report", + labels={}, + args=[], + kwargs={}, + cron="* * * * *", + ) + nested_task = root_task.model_copy(update={"schedule_id": "nested_task"}) + await schedule_source.add_schedule(root_task) + await nested.add_schedule(nested_task) + with caplog.at_level(logging.WARNING, logger="taskiq_nats.schedule_source"): + assert await schedule_source.get_schedules() == [root_task] + assert f"{nested.prefix}.nested_task" in caplog.text + assert await nested.get_schedules() == [nested_task] + finally: + await nested.shutdown() + + async def test_invalid_payload_still_raises_and_recovers( schedule_source: NATSKeyValueScheduleSource, ) -> None: @@ -129,3 +163,25 @@ async def test_invalid_payload_still_raises_and_recovers( await schedule_source.get_schedules() await schedule_source.kv.delete(key) assert await schedule_source.get_schedules() == [] + + +async def test_undecodable_multi_token_key_is_skipped( + schedule_source: NATSKeyValueScheduleSource, + caplog: pytest.LogCaptureFixture, +) -> None: + # Such keys were never read before multi-token IDs were supported, and + # they cannot be told apart from a nested source's schedules. + task = ScheduledTask( + schedule_id="valid", + task_name="report", + labels={}, + args=[], + kwargs={}, + cron="* * * * *", + ) + await schedule_source.add_schedule(task) + key = f"{schedule_source.prefix}.daily.broken" + await schedule_source.kv.put(key, b"not-pickle") + with caplog.at_level(logging.WARNING, logger="taskiq_nats.schedule_source"): + assert await schedule_source.get_schedules() == [task] + assert key in caplog.text From 49ad0389c97ba928d32e7ec9ef5a598a22d294c5 Mon Sep 17 00:00:00 2001 From: Wei Chen <182598863+cw-jlu@users.noreply.github.com> Date: Fri, 13 Nov 2026 22:31:48 +0800 Subject: [PATCH 3/3] Revert "Keep nested-prefix payloads from breaking the snapshot" This reverts commit da0f9c5b7da4241c888271f4e667f859c19014e4. Undecodable payloads under the prefix raise again. Sources sharing a bucket with nested prefixes need compatible serializers. Co-Authored-By: Claude Opus 5.5 --- taskiq_nats/schedule_source.py | 19 ++---------- tests/test_schedule_ids.py | 56 ---------------------------------- 2 files changed, 3 insertions(+), 72 deletions(-) diff --git a/taskiq_nats/schedule_source.py b/taskiq_nats/schedule_source.py index d7e11de..5e26a00 100644 --- a/taskiq_nats/schedule_source.py +++ b/taskiq_nats/schedule_source.py @@ -103,22 +103,9 @@ async def get_schedules(self) -> list[ScheduledTask]: result = [] for schedule in schedules: if schedule and schedule.value: - try: - task = ScheduledTask.model_validate( - self.serializer.loadb(schedule.value), - ) - except Exception: - # A multi-token key may belong to a source with a nested - # prefix and a different serializer, so it must not break - # this snapshot. Single-token keys can only be ours. - if "." not in schedule.key[len(self.prefix) + 1 :]: - raise - log.warning( - "Skipping undecodable schedule %r: it may belong to " - "a source with a nested prefix.", - schedule.key, - ) - continue + task = ScheduledTask.model_validate( + self.serializer.loadb(schedule.value), + ) # A nested prefix can match the wildcard without belonging to us. if schedule.key == f"{self.prefix}.{task.schedule_id}": result.append(task) diff --git a/tests/test_schedule_ids.py b/tests/test_schedule_ids.py index 85e05ad..fb5c267 100644 --- a/tests/test_schedule_ids.py +++ b/tests/test_schedule_ids.py @@ -1,5 +1,4 @@ import datetime as dt -import logging import pickle import uuid from collections.abc import AsyncGenerator @@ -7,7 +6,6 @@ import pytest from taskiq import ScheduledTask from taskiq.cli.scheduler.run import get_schedules -from taskiq.serializers import JSONSerializer from taskiq_nats import NATSKeyValueScheduleSource @@ -122,38 +120,6 @@ async def test_nested_prefix_keeps_its_own_schedules( await nested.shutdown() -async def test_nested_prefix_with_another_serializer( - nats_urls: list[str], - schedule_source: NATSKeyValueScheduleSource, - caplog: pytest.LogCaptureFixture, -) -> None: - nested = NATSKeyValueScheduleSource( - servers=nats_urls, - bucket_name=schedule_source.bucket_name, - prefix=f"{schedule_source.prefix}.nested", - serializer=JSONSerializer(), - ) - await nested.startup() - try: - root_task = ScheduledTask( - schedule_id="root_task", - task_name="report", - labels={}, - args=[], - kwargs={}, - cron="* * * * *", - ) - nested_task = root_task.model_copy(update={"schedule_id": "nested_task"}) - await schedule_source.add_schedule(root_task) - await nested.add_schedule(nested_task) - with caplog.at_level(logging.WARNING, logger="taskiq_nats.schedule_source"): - assert await schedule_source.get_schedules() == [root_task] - assert f"{nested.prefix}.nested_task" in caplog.text - assert await nested.get_schedules() == [nested_task] - finally: - await nested.shutdown() - - async def test_invalid_payload_still_raises_and_recovers( schedule_source: NATSKeyValueScheduleSource, ) -> None: @@ -163,25 +129,3 @@ async def test_invalid_payload_still_raises_and_recovers( await schedule_source.get_schedules() await schedule_source.kv.delete(key) assert await schedule_source.get_schedules() == [] - - -async def test_undecodable_multi_token_key_is_skipped( - schedule_source: NATSKeyValueScheduleSource, - caplog: pytest.LogCaptureFixture, -) -> None: - # Such keys were never read before multi-token IDs were supported, and - # they cannot be told apart from a nested source's schedules. - task = ScheduledTask( - schedule_id="valid", - task_name="report", - labels={}, - args=[], - kwargs={}, - cron="* * * * *", - ) - await schedule_source.add_schedule(task) - key = f"{schedule_source.prefix}.daily.broken" - await schedule_source.kv.put(key, b"not-pickle") - with caplog.at_level(logging.WARNING, logger="taskiq_nats.schedule_source"): - assert await schedule_source.get_schedules() == [task] - assert key in caplog.text