omit_exception for generator methods sscan_iter and iter_keys

Author: amirreza-sf80

django_redis/cache.py

@@ -13,16 +13,15 @@
 
 
 def omit_exception(
-    method: Optional[Callable] = None,
-    return_value: Optional[Any] = None,
+    method: Optional[Callable] = None, return_value: Optional[Any] = None, gen=False
 ):
     """
     Simple decorator that intercepts connection
     errors and ignores these if settings specify this.
     """
 
     if method is None:
-        return functools.partial(omit_exception, return_value=return_value)
+        return functools.partial(omit_exception, return_value=return_value, gen=gen)
 
     @functools.wraps(method)
     def _decorator(self, *args, **kwargs):
@@ -36,7 +35,21 @@ def _decorator(self, *args, **kwargs):
                 return return_value
             raise e.__cause__  # noqa: B904
 
-    return _decorator
+    @functools.wraps(method)
+    def _generator_decorator(self, *args, **kwargs):
+        try:
+            yield from method(self, *args, **kwargs)
+        except ConnectionInterrupted as e:
+            if self._ignore_exceptions:
+                if self._log_ignored_exceptions:
+                    self.logger.exception("Exception ignored")
+
+                return return_value
+            raise e.__cause__  # noqa: B904
+
+    if not gen:
+        return _decorator
+    return _generator_decorator
 
 
 class RedisCache(BaseCache):
@@ -147,7 +160,7 @@ def has_key(self, *args, **kwargs):
     def keys(self, *args, **kwargs):
         return self.client.keys(*args, **kwargs)
 
-    @omit_exception
+    @omit_exception(gen=True)
     def iter_keys(self, *args, **kwargs):
         return self.client.iter_keys(*args, **kwargs)
 
@@ -243,7 +256,7 @@ def srem(self, *args, **kwargs):
     def sscan(self, *args, **kwargs):
         return self.client.sscan(*args, **kwargs)
 
-    @omit_exception
+    @omit_exception(gen=True)
     def sscan_iter(self, *args, **kwargs):
         return self.client.sscan_iter(*args, **kwargs)
 

django_redis/client/default.py

@@ -775,8 +775,11 @@ def iter_keys(
             client = self.get_client(write=False)
 
         pattern = self.make_pattern(search, version=version)
-        for item in client.scan_iter(match=pattern, count=itersize):
-            yield self.reverse_key(item.decode())
+        try:
+            for item in client.scan_iter(match=pattern, count=itersize):
+                yield self.reverse_key(item.decode())
+        except _main_exceptions as e:
+            raise ConnectionInterrupted(connection=client) from e
 
     def keys(
         self,
@@ -1054,12 +1057,15 @@ def sscan_iter(
             client = self.get_client(write=False)
 
         key = self.make_key(key, version=version)
-        for value in client.sscan_iter(
-            key,
-            match=cast("PatternT", self.encode(match)) if match else None,
-            count=count,
-        ):
-            yield self.decode(value)
+        try:
+            for value in client.sscan_iter(
+                key,
+                match=cast("PatternT", self.encode(match)) if match else None,
+                count=count,
+            ):
+                yield self.decode(value)
+        except _main_exceptions as e:
+            raise ConnectionInterrupted(connection=client) from e
 
     def sunion(
         self,

tests/test_cache_options.py

@@ -3,13 +3,18 @@
 from typing import cast
 
 import pytest
-from django.core.cache import caches
+from django.core.cache import caches, cache as default_cache
 from pytest import LogCaptureFixture
 from redis.exceptions import ConnectionError as RedisConnectionError
 
 from django_redis.cache import RedisCache
 from django_redis.client import ShardClient
 
+iter_methods = {
+    "iter_keys",
+    "sscan_iter",
+}
+
 
 def make_key(key: str, prefix: str, version: str) -> str:
     return f"{prefix}#{version}#{key}"
@@ -55,6 +60,16 @@ def test_get_django_omit_exceptions(
     )
 
 
+def test_iterator_methods(ignore_exceptions_cache: RedisCache, subtests):
+    for m in iter_methods:
+        method = getattr(ignore_exceptions_cache, m)
+        with subtests.test(method=method):
+            if isinstance(default_cache.client, ShardClient) and m == "iter_keys":
+                pytest.skip(f"shard client doesn't support {m}")
+            for _ in method("abc"):
+                pass
+
+
 def test_get_django_omit_exceptions_priority_1(settings):
     caches_setting = copy.deepcopy(settings.CACHES)
     caches_setting["doesnotexist"]["OPTIONS"]["IGNORE_EXCEPTIONS"] = True