Added serializers and fixed tests.

Signed-off-by: Pavel Kirilin <win10@list.ru>

Author: s3rius

poetry.lock

@@ -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

pyproject.toml

@@ -1,4 +1,3 @@
-import pickle
 import sys
 from contextlib import asynccontextmanager
 from typing import (
@@ -15,16 +14,18 @@
 
 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 @@
 
 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 @@ def __init__(
  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 @@ def __init__(
  :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 @@ async def set_result(
  :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 @@ async def get_result(
  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 @@ def __init__(
  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 @@ def __init__(
  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 @@ async def set_result(
  """
  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 @@ async def get_result(
  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 @@ def __init__(
  **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 @@ async def set_result(
  """
  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 @@ async def get_result(
  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]

taskiq_redis/redis_backend.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 @@
 
 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 @@ def __init__(
  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,

taskiq_redis/redis_broker.py

@@ -4,7 +4,7 @@
 
 from redis.asyncio import (
  BlockingConnectionPool,
- ConnectionPool,
+ Connection,
  Redis,
  RedisCluster,
  Sentinel,
@@ -13,8 +13,7 @@
 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 @@
 
 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 @@ def __init__(
  **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 @@ async def post_send(self, task: ScheduledTask) -> None:
  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 @@ async def post_send(self, task: ScheduledTask) -> None:
  """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]

taskiq_redis/schedule_source.py

@@ -90,34 +90,29 @@ 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,
 ) -> None:
  """
  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.
  """
  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()