From e95c77e59316d067da414c697a826c856f2446fc Mon Sep 17 00:00:00 2001 From: w1nstep <3539832429@qq.com> Date: Wed, 2 Sep 2026 10:55:49 +0800 Subject: [PATCH] fix(redis): make legacy key cleanup safer --- scripts/cleanup_legacy_rate_limit_keys.py | 33 +++++++---- .../test_cleanup_legacy_rate_limit_keys.py | 59 +++++++++++++++++-- 2 files changed, 76 insertions(+), 16 deletions(-) diff --git a/scripts/cleanup_legacy_rate_limit_keys.py b/scripts/cleanup_legacy_rate_limit_keys.py index 4399e5d08d7..c232ffc6d17 100644 --- a/scripts/cleanup_legacy_rate_limit_keys.py +++ b/scripts/cleanup_legacy_rate_limit_keys.py @@ -10,15 +10,18 @@ Usage: python scripts/cleanup_legacy_rate_limit_keys.py [options] The Redis connection is read from LiteLLM's normal REDIS_* configuration. A -namespace can be supplied to narrow the SCAN. redis-py RedisCluster's -``scan_iter`` is used as-is, so cluster scans remain node-aware. Deletions are -sent one key at a time to avoid a cross-slot multi-key command. +namespace can be supplied to narrow the SCAN. Apply mode requires a namespace +and removes only keys with a permanent TTL (``TTL=-1``), so a dry-run is the +only unscoped mode and active finite-TTL counters are retained. redis-py +RedisCluster's ``scan_iter`` is used as-is, so cluster scans remain node-aware. +Deletions are sent one key at a time to avoid a cross-slot multi-key command. """ from __future__ import annotations import argparse import json +import os import re import sys from collections import Counter @@ -128,6 +131,8 @@ def cleanup_legacy_keys( normalized_namespace = _normalize_namespace(namespace) if count <= 0: raise ValueError("count must be positive") + if apply and normalized_namespace is None: + raise ValueError("namespace is required when apply=True") scan_iter = getattr(client, "scan_iter", None) ttl = getattr(client, "ttl", None) @@ -153,14 +158,16 @@ def cleanup_legacy_keys( report.ttl_errors += 1 continue + is_permanent = False if ttl_value == -1: report.permanent += 1 + is_permanent = True elif isinstance(ttl_value, int) and ttl_value >= 0: report.finite += 1 else: report.unknown_ttl += 1 - if apply: + if apply and is_permanent: try: _delete_one(client, key) except Exception: @@ -172,8 +179,12 @@ def cleanup_legacy_keys( def _connection_kwargs(args: argparse.Namespace) -> dict[str, object]: + if args.host is None and (args.port is not None or args.db is not None): + raise ValueError("--host is required when using --port or --db") + if args.host is not None and args.port is None: + raise ValueError("--port is required when using --host") + values = { - "url": args.url, "host": args.host, "port": args.port, "db": args.db, @@ -187,7 +198,10 @@ def _connect(args: argparse.Namespace) -> object: # and Cluster/Sentinel connection handling for the real command. from litellm._redis import get_redis_client - return get_redis_client(**_connection_kwargs(args)) + connection_kwargs = _connection_kwargs(args) + if connection_kwargs and (os.getenv("REDIS_CLUSTER_NODES") or os.getenv("REDIS_SENTINEL_NODES")): + raise ValueError("explicit host overrides cannot be combined with Redis Cluster or Sentinel configuration") + return get_redis_client(**connection_kwargs) def _write_report(report: CleanupReport, *, apply: bool, json_output: bool) -> None: @@ -216,10 +230,9 @@ def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--namespace", help="only inspect keys beginning with NAMESPACE:") parser.add_argument("--count", type=int, default=1000, help="Redis SCAN count hint (default: 1000)") - parser.add_argument("--url", help="Redis URL; otherwise use the normal REDIS_* configuration") - parser.add_argument("--host", help="Redis host override") - parser.add_argument("--port", type=int, help="Redis port override") - parser.add_argument("--db", type=int, help="Redis database override") + parser.add_argument("--host", help="Redis host override; use with --port") + parser.add_argument("--port", type=int, help="Redis port override; use with --host") + parser.add_argument("--db", type=int, help="Redis database override; use with --host") parser.add_argument("--json", action="store_true", dest="json_output", help="emit machine-readable JSON") parser.add_argument( "--apply", diff --git a/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py b/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py index f7291d47879..0b5425d26f4 100644 --- a/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py +++ b/tests/test_litellm/test_cleanup_legacy_rate_limit_keys.py @@ -2,6 +2,7 @@ import importlib.util import sys +from argparse import Namespace from pathlib import Path import pytest @@ -84,28 +85,74 @@ def test_dry_run_scans_cluster_aware_iterator_without_deleting() -> None: assert redis.unlink_calls == [] -def test_apply_deletes_only_unique_legacy_keys() -> None: +def test_apply_requires_namespace() -> None: old_key = "global_router:id:model:rpm:13-07" redis = FakeRedis([old_key, old_key, "global_router:id:model:rpm:v2:29300000"], {old_key: -1}) - report = cleanup.cleanup_legacy_keys(redis, apply=True) + with pytest.raises(ValueError, match="namespace"): + cleanup.cleanup_legacy_keys(redis, apply=True) - assert report.candidates == 1 + assert redis.unlink_calls == [] + + +def test_apply_deletes_only_unique_permanent_legacy_keys() -> None: + permanent_key = "tenant-a:global_router:id:model:rpm:13-07" + finite_key = "tenant-a:global_router:id:model:tpm:13-07" + redis = FakeRedis( + [permanent_key, permanent_key, finite_key, "tenant-a:global_router:id:model:rpm:v2:29300000"], + {permanent_key: -1, finite_key: 30}, + ) + + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", apply=True) + + assert report.candidates == 2 + assert report.permanent == 1 + assert report.finite == 1 assert report.deleted == 1 assert report.failed == 0 - assert redis.unlink_calls == [old_key] + assert redis.unlink_calls == [permanent_key] def test_apply_falls_back_to_single_key_delete() -> None: - old_key = "13-07:model" + old_key = "tenant-a:13-07:model" redis = DeleteOnlyRedis([old_key], {old_key: 9}) - report = cleanup.cleanup_legacy_keys(redis, apply=True) + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", apply=True) + + assert report.deleted == 0 + assert redis.delete_calls == [] + + +def test_apply_falls_back_to_single_key_delete_for_permanent_key() -> None: + old_key = "tenant-a:13-07:model" + redis = DeleteOnlyRedis([old_key], {old_key: -1}) + + report = cleanup.cleanup_legacy_keys(redis, namespace="tenant-a", apply=True) assert report.deleted == 1 assert redis.delete_calls == [old_key] +def test_connection_overrides_require_an_explicit_host_and_port() -> None: + with pytest.raises(ValueError, match="--host"): + cleanup._connection_kwargs(Namespace(host=None, port=6380, db=None)) + with pytest.raises(ValueError, match="--port"): + cleanup._connection_kwargs(Namespace(host="redis.example", port=None, db=None)) + + +def test_connection_overrides_do_not_accept_a_url() -> None: + kwargs = cleanup._connection_kwargs(Namespace(host="redis.example", port=6380, db=2)) + + assert kwargs == {"host": "redis.example", "port": 6380, "db": 2} + + +def test_connection_overrides_do_not_mix_with_cluster_configuration(monkeypatch) -> None: + monkeypatch.setenv("REDIS_CLUSTER_NODES", '[{"host": "cluster.example", "port": 6379}]') + + with pytest.raises(ValueError, match="Cluster or Sentinel"): + cleanup._connect(Namespace(host="redis.example", port=6380, db=None)) + + def test_count_must_be_positive() -> None: with pytest.raises(ValueError, match="positive"): cleanup.cleanup_legacy_keys(FakeRedis([], {}), count=0)