Skip to content

Commit cddce60

Browse files
committed
omit_exception for generator methods sscan_iter and iter_keys
1 parent 755dfb2 commit cddce60

File tree

3 files changed

+49
-12
lines changed

3 files changed

+49
-12
lines changed

django_redis/cache.py

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,15 @@
1515
def omit_exception(
1616
method: Optional[Callable] = None,
1717
return_value: Optional[Any] = None,
18+
gen=False,
1819
):
1920
"""
2021
Simple decorator that intercepts connection
2122
errors and ignores these if settings specify this.
2223
"""
2324

2425
if method is None:
25-
return functools.partial(omit_exception, return_value=return_value)
26+
return functools.partial(omit_exception, return_value=return_value, gen=gen)
2627

2728
@functools.wraps(method)
2829
def _decorator(self, *args, **kwargs):
@@ -36,7 +37,21 @@ def _decorator(self, *args, **kwargs):
3637
return return_value
3738
raise e.__cause__ # noqa: B904
3839

39-
return _decorator
40+
@functools.wraps(method)
41+
def _generator_decorator(self, *args, **kwargs):
42+
try:
43+
yield from method(self, *args, **kwargs)
44+
except ConnectionInterrupted as e:
45+
if self._ignore_exceptions:
46+
if self._log_ignored_exceptions:
47+
self.logger.exception("Exception ignored")
48+
49+
return return_value
50+
raise e.__cause__ # noqa: B904
51+
52+
if not gen:
53+
return _decorator
54+
return _generator_decorator
4055

4156

4257
class RedisCache(BaseCache):
@@ -147,7 +162,7 @@ def has_key(self, *args, **kwargs):
147162
def keys(self, *args, **kwargs):
148163
return self.client.keys(*args, **kwargs)
149164

150-
@omit_exception
165+
@omit_exception(gen=True)
151166
def iter_keys(self, *args, **kwargs):
152167
return self.client.iter_keys(*args, **kwargs)
153168

@@ -243,7 +258,7 @@ def srem(self, *args, **kwargs):
243258
def sscan(self, *args, **kwargs):
244259
return self.client.sscan(*args, **kwargs)
245260

246-
@omit_exception
261+
@omit_exception(gen=True)
247262
def sscan_iter(self, *args, **kwargs):
248263
return self.client.sscan_iter(*args, **kwargs)
249264

django_redis/client/default.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -775,8 +775,11 @@ def iter_keys(
775775
client = self.get_client(write=False)
776776

777777
pattern = self.make_pattern(search, version=version)
778-
for item in client.scan_iter(match=pattern, count=itersize):
779-
yield self.reverse_key(item.decode())
778+
try:
779+
for item in client.scan_iter(match=pattern, count=itersize):
780+
yield self.reverse_key(item.decode())
781+
except _main_exceptions as e:
782+
raise ConnectionInterrupted(connection=client) from e
780783

781784
def keys(
782785
self,
@@ -1054,12 +1057,15 @@ def sscan_iter(
10541057
client = self.get_client(write=False)
10551058

10561059
key = self.make_key(key, version=version)
1057-
for value in client.sscan_iter(
1058-
key,
1059-
match=cast("PatternT", self.encode(match)) if match else None,
1060-
count=count,
1061-
):
1062-
yield self.decode(value)
1060+
try:
1061+
for value in client.sscan_iter(
1062+
key,
1063+
match=cast("PatternT", self.encode(match)) if match else None,
1064+
count=count,
1065+
):
1066+
yield self.decode(value)
1067+
except _main_exceptions as e:
1068+
raise ConnectionInterrupted(connection=client) from e
10631069

10641070
def sunion(
10651071
self,

tests/test_cache_options.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,19 @@
33
from typing import cast
44

55
import pytest
6+
from django.core.cache import cache as default_cache
67
from django.core.cache import caches
78
from pytest import LogCaptureFixture
89
from redis.exceptions import ConnectionError as RedisConnectionError
910

1011
from django_redis.cache import RedisCache
1112
from django_redis.client import ShardClient
1213

14+
iter_methods = {
15+
"iter_keys",
16+
"sscan_iter",
17+
}
18+
1319

1420
def make_key(key: str, prefix: str, version: str) -> str:
1521
return f"{prefix}#{version}#{key}"
@@ -55,6 +61,16 @@ def test_get_django_omit_exceptions(
5561
)
5662

5763

64+
def test_iterator_methods(ignore_exceptions_cache: RedisCache, subtests):
65+
for m in iter_methods:
66+
method = getattr(ignore_exceptions_cache, m)
67+
with subtests.test(method=method):
68+
if isinstance(default_cache.client, ShardClient) and m == "iter_keys":
69+
pytest.skip(f"shard client doesn't support {m}")
70+
for _ in method("abc"):
71+
pass
72+
73+
5874
def test_get_django_omit_exceptions_priority_1(settings):
5975
caches_setting = copy.deepcopy(settings.CACHES)
6076
caches_setting["doesnotexist"]["OPTIONS"]["IGNORE_EXCEPTIONS"] = True

0 commit comments

Comments
 (0)