Repository: taskiq-python/taskiq-redis

Issue:
Custom serializers
Would be great to have custom serializers in set/get results, or at least JSON, like it is implemented in the main project TaskIQ.

Hints from issue comments:
@s3rius: True. Will add soon.@s3rius: Look, I've checked the code and you can do so already. You can set custom serializers in the constructor. @quizmoon: @s3rius , in the backend part it is hardcoded, e.g in "set_result":

```
        redis_set_params: Dict[str, Union[str, bytes, int]] = {
            "name": task_id,
            "value": pickle.dumps(result),
        }
```@s3rius: I see now. Accidentally skipped it. We will add a new parameter to the constructor.

Solution patch:
diff --git a/pyproject.toml b/pyproject.toml
index d5155ed..9dc59a9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -26,7 +26,7 @@ keywords = [
 
 [tool.poetry.dependencies]
 python = "^3.8.1"
-taskiq = ">=0.10.3,<1"
+taskiq = ">=0.11.1,<1"
 redis = "^5"
 
 [tool.poetry.group.dev.dependencies]
@@ -40,7 +40,7 @@ fakeredis = "^2"
 pre-commit = "^2.20.0"
 pytest-xdist = { version = "^2.5.0", extras = ["psutil"] }
 ruff = "^0.1.0"
-types-redis = "^4.6.0.7"
+types-redis = "^4.6.0.20240425"
 
 [tool.mypy]
 strict = true
diff --git a/taskiq_redis/redis_backend.py b/taskiq_redis/redis_backend.py
index dbd9e20..28ce927 100644
--- a/taskiq_redis/redis_backend.py
+++ b/taskiq_redis/redis_backend.py
@@ -1,4 +1,3 @@
-import pickle
 import sys
 from contextlib import asynccontextmanager
 from typing import (
@@ -15,16 +14,18 @@ from typing import (
 
 from redis.asyncio import BlockingConnectionPool, Redis, Sentinel
 from redis.asyncio.cluster import RedisCluster
+from redis.asyncio.connection import Connection
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
 from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.compat import model_dump, model_validate
+from taskiq.serializers import PickleSerializer
 
 from taskiq_redis.exceptions import (
     DuplicateExpireTimeSelectedError,
     ExpireTimeMustBeMoreThanZeroError,
     ResultIsMissingError,
 )
-from taskiq_redis.serializer import PickleSerializer
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -33,8 +34,10 @@ else:
 
 if TYPE_CHECKING:
     _Redis: TypeAlias = Redis[bytes]
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
 else:
     _Redis: TypeAlias = Redis
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
 
 _ReturnType = TypeVar("_ReturnType")
 
@@ -49,6 +52,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
         max_connection_pool_size: Optional[int] = None,
+        serializer: Optional[TaskiqSerializer] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -66,11 +70,12 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
             and result_px_time are equal zero.
         """
-        self.redis_pool = BlockingConnectionPool.from_url(
+        self.redis_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
             url=redis_url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
         )
+        self.serializer = serializer or PickleSerializer()
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -110,9 +115,9 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :param task_id: ID of the task.
         :param result: TaskiqResult instance.
         """
-        redis_set_params: Dict[str, Union[str, bytes, int]] = {
+        redis_set_params: Dict[str, Union[str, int, bytes]] = {
             "name": task_id,
-            "value": pickle.dumps(result),
+            "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
             redis_set_params["ex"] = self.result_ex_time
@@ -159,8 +164,9 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         if result_value is None:
             raise ResultIsMissingError
 
-        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(  # noqa: S301
-            result_value,
+        taskiq_result = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_value),
         )
 
         if not with_logs:
@@ -178,6 +184,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         keep_results: bool = True,
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
+        serializer: Optional[TaskiqSerializer] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -198,6 +205,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
             redis_url,
             **connection_kwargs,
         )
+        self.serializer = serializer or PickleSerializer()
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -239,7 +247,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         """
         redis_set_params: Dict[str, Union[str, bytes, int]] = {
             "name": task_id,
-            "value": pickle.dumps(result),
+            "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
             redis_set_params["ex"] = self.result_ex_time
@@ -283,8 +291,9 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         if result_value is None:
             raise ResultIsMissingError
 
-        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(  # noqa: S301
-            result_value,
+        taskiq_result: TaskiqResult[_ReturnType] = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_value),
         )
 
         if not with_logs:
@@ -331,9 +340,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
             **connection_kwargs,
         )
         self.master_name = master_name
-        if serializer is None:
-            serializer = PickleSerializer()
-        self.serializer = serializer
+        self.serializer = serializer or PickleSerializer()
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -375,7 +382,7 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         """
         redis_set_params: Dict[str, Union[str, bytes, int]] = {
             "name": task_id,
-            "value": self.serializer.dumpb(result),
+            "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
             redis_set_params["ex"] = self.result_ex_time
@@ -422,11 +429,17 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
         if result_value is None:
             raise ResultIsMissingError
 
-        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(  # noqa: S301
-            result_value,
+        taskiq_result = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_value),
         )
 
         if not with_logs:
             taskiq_result.log = None
 
         return taskiq_result
+
+    async def shutdown(self) -> None:
+        """Shutdown sentinel connections."""
+        for sentinel in self.sentinel.sentinels:
+            await sentinel.aclose()  # type: ignore[attr-defined]
diff --git a/taskiq_redis/redis_broker.py b/taskiq_redis/redis_broker.py
index 2fab345..ec2c13c 100644
--- a/taskiq_redis/redis_broker.py
+++ b/taskiq_redis/redis_broker.py
@@ -1,7 +1,8 @@
+import sys
 from logging import getLogger
-from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Optional, TypeVar
 
-from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, Connection, Redis
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.message import BrokerMessage
@@ -10,6 +11,16 @@ _T = TypeVar("_T")
 
 logger = getLogger("taskiq.redis_broker")
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+else:
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
+
 
 class BaseRedisBroker(AsyncBroker):
     """Base broker that works with Redis."""
@@ -40,7 +51,7 @@ class BaseRedisBroker(AsyncBroker):
             task_id_generator=task_id_generator,
         )
 
-        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
+        self.connection_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
diff --git a/taskiq_redis/schedule_source.py b/taskiq_redis/schedule_source.py
index 6043e1f..a6d5b54 100644
--- a/taskiq_redis/schedule_source.py
+++ b/taskiq_redis/schedule_source.py
@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, List, Optional, Tuple
 
 from redis.asyncio import (
     BlockingConnectionPool,
-    ConnectionPool,
+    Connection,
     Redis,
     RedisCluster,
     Sentinel,
@@ -13,8 +13,7 @@ from taskiq import ScheduleSource
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
 from taskiq.scheduler.scheduled_task import ScheduledTask
-
-from taskiq_redis.serializer import PickleSerializer
+from taskiq.serializers import PickleSerializer
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -23,8 +22,10 @@ else:
 
 if TYPE_CHECKING:
     _Redis: TypeAlias = Redis[bytes]
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
 else:
     _Redis: TypeAlias = Redis
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
 
 
 class RedisScheduleSource(ScheduleSource):
@@ -53,7 +54,7 @@ class RedisScheduleSource(ScheduleSource):
         **connection_kwargs: Any,
     ) -> None:
         self.prefix = prefix
-        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
+        self.connection_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
@@ -186,6 +187,10 @@ class RedisClusterScheduleSource(ScheduleSource):
         if task.time is not None:
             await self.delete_schedule(task.schedule_id)
 
+    async def shutdown(self) -> None:
+        """Shut down the schedule source."""
+        await self.redis.aclose()  # type: ignore[attr-defined]
+
 
 class RedisSentinelScheduleSource(ScheduleSource):
     """
@@ -279,3 +284,8 @@ class RedisSentinelScheduleSource(ScheduleSource):
         """Delete a task after it's completed."""
         if task.time is not None:
             await self.delete_schedule(task.schedule_id)
+
+    async def shutdown(self) -> None:
+        """Shut down the schedule source."""
+        for sentinel in self.sentinel.sentinels:
+            await sentinel.aclose()  # type: ignore[attr-defined]
diff --git a/taskiq_redis/serializer.py b/taskiq_redis/serializer.py
deleted file mode 100644
index b5bec92..0000000
--- a/taskiq_redis/serializer.py
+++ /dev/null
@@ -1,16 +0,0 @@
-import pickle
-from typing import Any
-
-from taskiq.abc.serializer import TaskiqSerializer
-
-
-class PickleSerializer(TaskiqSerializer):
-    """Serializer that uses pickle."""
-
-    def dumpb(self, value: Any) -> bytes:
-        """Dumps value to bytes."""
-        return pickle.dumps(value)
-
-    def loadb(self, value: bytes) -> Any:
-        """Loads value from bytes."""
-        return pickle.loads(value)  # noqa: S301


Test patch:
diff --git a/tests/test_backend.py b/tests/test_backend.py
index 6946e49..aad0075 100644
--- a/tests/test_backend.py
+++ b/tests/test_backend.py
@@ -90,7 +90,7 @@ async def test_success_backend_default_result(
 
 
 @pytest.mark.anyio
-async def test_success_backend_custom_result(
+async def test_error_backend_custom_result(
     custom_taskiq_result: TaskiqResult[_ReturnType],
     task_id: str,
     redis_url: str,
@@ -98,6 +98,8 @@ async def test_success_backend_custom_result(
     """
     Tests normal behavior with custom result in TaskiqResult.
 
+    Setting custom class as a result should raise an error.
+
     :param custom_taskiq_result: TaskiqResult with custom result.
     :param task_id: ID for task.
     :param redis_url: url to redis.
@@ -105,19 +107,12 @@ async def test_success_backend_custom_result(
     backend: RedisAsyncResultBackend[_ReturnType] = RedisAsyncResultBackend(
         redis_url,
     )
-    await backend.set_result(
-        task_id=task_id,
-        result=custom_taskiq_result,
-    )
-    result = await backend.get_result(task_id=task_id)
+    with pytest.raises(ValueError):
+        await backend.set_result(
+            task_id=task_id,
+            result=custom_taskiq_result,
+        )
 
-    assert (
-        result.return_value.test_arg  # type: ignore
-        == custom_taskiq_result.return_value.test_arg  # type: ignore
-    )
-    assert result.is_err == custom_taskiq_result.is_err
-    assert result.execution_time == custom_taskiq_result.execution_time
-    assert result.log == custom_taskiq_result.log
     await backend.shutdown()