taskiq-redis 0.1.0__tar.gz → 0.2.1__tar.gz

{taskiq_redis-0.1.0 → taskiq_redis-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: taskiq-redis
-Version: 0.1.0
+Version: 0.2.1
 Summary: Redis integration for taskiq
 Home-page: https://github.com/taskiq-python/taskiq-redis
 Keywords: taskiq,tasks,distributed,async,redis,result_backend
@@ -43,18 +43,19 @@ pip install taskiq-redis
 # Usage
 
 Let's see the example with the redis broker and redis async result:
+
 ```python
 import asyncio
 
-from taskiq_redis.redis_broker import RedisBroker
+from taskiq_redis.redis_broker import ListQueueBroker
 from taskiq_redis.redis_backend import RedisAsyncResultBackend
 
-
 redis_async_result = RedisAsyncResultBackend(
     redis_url="redis://localhost:6379",
 )
 
-broker = RedisBroker(
+# Or you can use PubSubBroker if you need broadcasting
+broker = ListQueueBroker(
     url="redis://localhost:6379",
     result_backend=redis_async_result,
 )
@@ -75,9 +76,17 @@ async def main():
 asyncio.run(main())
 ```
 
-## RedisBroker configuration
+## PubSubBroker and ListQueueBroker configuration
+
+We have two brokers with similar interfaces, but with different logic.
+The PubSubBroker uses redis' pubsub mechanism and is very powerful,
+but it executes every task on all workers, because PUBSUB broadcasts message
+to all subscribers.
 
-RedisBroker parameters:
+If you want your messages to be processed only once, please use ListQueueBroker.
+It uses redis' [LPUSH](https://redis.io/commands/lpush/) and [BRPOP](https://redis.io/commands/brpop/) commands to deal with messages.
+
+Brokers parameters:
 * `url` - url to redis.
 * `task_id_generator` - custom task_id genertaor.
 * `result_backend` - custom result backend.
@@ -89,4 +98,19 @@ RedisBroker parameters:
 RedisAsyncResultBackend parameters:
 * `redis_url` - url to redis.
 * `keep_results` - flag to not remove results from Redis after reading.
-
+* `result_ex_time` - expire time in seconds (by default - 1 minute)
+* `result_px_time` - expire time in milliseconds (by default - not specified)
+> IMPORTANT: You must specify either `result_ex_time` or `result_px_time`.  
+>```python
+># First variant
+>redis_async_result = RedisAsyncResultBackend(
+>    redis_url="redis://localhost:6379",
+>    result_ex_time=1000,
+>)
+>
+># Second variant
+>redis_async_result = RedisAsyncResultBackend(
+>    redis_url="redis://localhost:6379",
+>    result_px_time=1000000,
+>)
+>```

{taskiq_redis-0.1.0 → taskiq_redis-0.2.1}/README.md

@@ -16,18 +16,19 @@ pip install taskiq-redis
 # Usage
 
 Let's see the example with the redis broker and redis async result:
+
 ```python
 import asyncio
 
-from taskiq_redis.redis_broker import RedisBroker
+from taskiq_redis.redis_broker import ListQueueBroker
 from taskiq_redis.redis_backend import RedisAsyncResultBackend
 
-
 redis_async_result = RedisAsyncResultBackend(
     redis_url="redis://localhost:6379",
 )
 
-broker = RedisBroker(
+# Or you can use PubSubBroker if you need broadcasting
+broker = ListQueueBroker(
     url="redis://localhost:6379",
     result_backend=redis_async_result,
 )
@@ -48,9 +49,17 @@ async def main():
 asyncio.run(main())
 ```
 
-## RedisBroker configuration
+## PubSubBroker and ListQueueBroker configuration
+
+We have two brokers with similar interfaces, but with different logic.
+The PubSubBroker uses redis' pubsub mechanism and is very powerful,
+but it executes every task on all workers, because PUBSUB broadcasts message
+to all subscribers.
+
+If you want your messages to be processed only once, please use ListQueueBroker.
+It uses redis' [LPUSH](https://redis.io/commands/lpush/) and [BRPOP](https://redis.io/commands/brpop/) commands to deal with messages.
 
-RedisBroker parameters:
+Brokers parameters:
 * `url` - url to redis.
 * `task_id_generator` - custom task_id genertaor.
 * `result_backend` - custom result backend.
@@ -62,3 +71,19 @@ RedisBroker parameters:
 RedisAsyncResultBackend parameters:
 * `redis_url` - url to redis.
 * `keep_results` - flag to not remove results from Redis after reading.
+* `result_ex_time` - expire time in seconds (by default - 1 minute)
+* `result_px_time` - expire time in milliseconds (by default - not specified)
+> IMPORTANT: You must specify either `result_ex_time` or `result_px_time`.  
+>```python
+># First variant
+>redis_async_result = RedisAsyncResultBackend(
+>    redis_url="redis://localhost:6379",
+>    result_ex_time=1000,
+>)
+>
+># Second variant
+>redis_async_result = RedisAsyncResultBackend(
+>    redis_url="redis://localhost:6379",
+>    result_px_time=1000000,
+>)
+>```

{taskiq_redis-0.1.0 → taskiq_redis-0.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-redis"
-version = "0.1.0"
+version = "0.2.1"
 description = "Redis integration for taskiq"
 authors = ["taskiq-team <taskiq@norely.com>"]
 readme = "README.md"

taskiq_redis-0.2.1/taskiq_redis/__init__.py

@@ -0,0 +1,5 @@
+"""Package for redis integration."""
+from taskiq_redis.redis_backend import RedisAsyncResultBackend
+from taskiq_redis.redis_broker import ListQueueBroker
+
+__all__ = ["RedisAsyncResultBackend", "ListQueueBroker"]

taskiq_redis-0.2.1/taskiq_redis/exceptions.py

@@ -0,0 +1,10 @@
+class TaskIQRedisError(Exception):
+    """Base error for all taskiq-redis exceptions."""
+
+
+class DuplicateExpireTimeSelectedError(TaskIQRedisError):
+    """Error if two lifetimes are selected."""
+
+
+class ExpireTimeMustBeMoreThanZeroError(TaskIQRedisError):
+    """Error if two lifetimes are less or equal zero."""

taskiq_redis-0.2.1/taskiq_redis/redis_backend.py

@@ -0,0 +1,126 @@
+import pickle
+from typing import Dict, Optional, TypeVar, Union
+
+from redis.asyncio import ConnectionPool, Redis
+from taskiq import AsyncResultBackend
+from taskiq.abc.result_backend import TaskiqResult
+
+from taskiq_redis.exceptions import (
+    DuplicateExpireTimeSelectedError,
+    ExpireTimeMustBeMoreThanZeroError,
+)
+
+_ReturnType = TypeVar("_ReturnType")
+
+
+class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
+    """Async result based on redis."""
+
+    def __init__(
+        self,
+        redis_url: str,
+        keep_results: bool = True,
+        result_ex_time: Optional[int] = None,
+        result_px_time: Optional[int] = None,
+    ):
+        """
+        Constructs a new result backend.
+
+        :param redis_url: url to redis.
+        :param keep_results: flag to not remove results from Redis after reading.
+        :param result_ex_time: expire time in seconds for result.
+        :param result_px_time: expire time in milliseconds for result.
+
+        :raises DuplicateExpireTimeSelectedError: if result_ex_time
+            and result_px_time are selected.
+        :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
+            and result_px_time are equal zero.
+        """
+        self.redis_pool = ConnectionPool.from_url(redis_url)
+        self.keep_results = keep_results
+        self.result_ex_time = result_ex_time
+        self.result_px_time = result_px_time
+
+        if self.result_ex_time == 0 or self.result_px_time == 0:
+            raise ExpireTimeMustBeMoreThanZeroError(
+                "You must select one expire time param and it must be more than zero.",
+            )
+
+        if self.result_ex_time and self.result_px_time:
+            raise DuplicateExpireTimeSelectedError(
+                "Choose either result_ex_time or result_px_time.",
+            )
+
+        if not self.result_ex_time and not self.result_px_time:
+            self.result_ex_time = 60
+
+    async def shutdown(self) -> None:
+        """Closes redis connection."""
+        await self.redis_pool.disconnect()
+        await super().shutdown()
+
+    async def set_result(
+        self,
+        task_id: str,
+        result: TaskiqResult[_ReturnType],
+    ) -> None:
+        """
+        Sets task result in redis.
+
+        Dumps TaskiqResult instance into the bytes and writes
+        it to redis.
+
+        :param task_id: ID of the task.
+        :param result: TaskiqResult instance.
+        """
+        redis_set_params: Dict[str, Union[str, bytes, int]] = {
+            "name": task_id,
+            "value": pickle.dumps(result),
+        }
+        if self.result_ex_time:
+            redis_set_params["ex"] = self.result_ex_time
+        elif self.result_px_time:
+            redis_set_params["px"] = self.result_px_time
+
+        async with Redis(connection_pool=self.redis_pool) as redis:
+            await redis.set(**redis_set_params)
+
+    async def is_result_ready(self, task_id: str) -> bool:
+        """
+        Returns whether the result is ready.
+
+        :param task_id: ID of the task.
+
+        :returns: True if the result is ready else False.
+        """
+        async with Redis(connection_pool=self.redis_pool) as redis:
+            return bool(await redis.exists(task_id))
+
+    async def get_result(  # noqa: WPS210
+        self,
+        task_id: str,
+        with_logs: bool = False,
+    ) -> TaskiqResult[_ReturnType]:
+        """
+        Gets result from the task.
+
+        :param task_id: task's id.
+        :param with_logs: if True it will download task's logs.
+        :return: task's return value.
+        """
+        async with Redis(connection_pool=self.redis_pool) as redis:
+            if self.keep_results:
+                result_value = await redis.get(
+                    name=task_id,
+                )
+            else:
+                result_value = await redis.getdel(
+                    name=task_id,
+                )
+
+        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(result_value)
+
+        if not with_logs:
+            taskiq_result.log = None
+
+        return taskiq_result

{taskiq_redis-0.1.0 → taskiq_redis-0.2.1}/taskiq_redis/redis_broker.py

@@ -1,4 +1,5 @@
 import pickle
+from abc import abstractmethod
 from logging import getLogger
 from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
@@ -12,8 +13,8 @@ _T = TypeVar("_T")  # noqa: WPS111
 logger = getLogger("taskiq.redis_broker")
 
 
-class RedisBroker(AsyncBroker):
-    """Broker that works with Redis."""
+class BaseRedisBroker(AsyncBroker):
+    """Base broker that works with Redis."""
 
     def __init__(
         self,
@@ -44,31 +45,12 @@ class RedisBroker(AsyncBroker):
             max_connections=max_connection_pool_size,
             **connection_kwargs,
         )
-
-        self.redis_pubsub_channel = queue_name
+        self.queue_name = queue_name
 
     async def shutdown(self) -> None:
         """Closes redis connection pool."""
         await self.connection_pool.disconnect()
 
-    async def kick(self, message: BrokerMessage) -> None:
-        """
-        Sends a message to the redis broker list.
-
-        This function constructs message for redis
-        and sends it.
-
-        The message is pickled dict object with message,
-        task_id, task_name and labels.
-
-        :param message: message to send.
-        """
-        async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.publish(
-                self.redis_pubsub_channel,
-                pickle.dumps(message),
-            )
-
     async def listen(self) -> AsyncGenerator[BrokerMessage, None]:
         """
         Listen redis queue for new messages.
@@ -78,24 +60,60 @@ class RedisBroker(AsyncBroker):
 
         :yields: broker messages.
         """
+        async for message in self._listen_to_raw_messages():
+            try:
+                redis_message = pickle.loads(message)
+                if isinstance(redis_message, BrokerMessage):
+                    yield redis_message
+            except (
+                TypeError,
+                AttributeError,
+                pickle.UnpicklingError,
+            ) as exc:
+                logger.debug(
+                    "Cannot read broker message %s",
+                    exc,
+                    exc_info=True,
+                )
+
+    @abstractmethod
+    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
+        """
+        Generator for reading raw data from Redis.
+
+        :yields: raw data.
+        """
+        yield  # type: ignore
+
+
+class PubSubBroker(BaseRedisBroker):
+    """Broker that works with Redis and broadcasts tasks to all workers."""
+
+    async def kick(self, message: BrokerMessage) -> None:  # noqa: D102
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            await redis_conn.publish(self.queue_name, pickle.dumps(message))
+
+    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
             redis_pubsub_channel = redis_conn.pubsub()
-            await redis_pubsub_channel.subscribe(self.redis_pubsub_channel)
+            await redis_pubsub_channel.subscribe(self.queue_name)
             async for message in redis_pubsub_channel.listen():
-                if message:
-                    try:
-                        redis_message = pickle.loads(
-                            message["data"],
-                        )
-                        if isinstance(redis_message, BrokerMessage):
-                            yield redis_message
-                    except (
-                        TypeError,
-                        AttributeError,
-                        pickle.UnpicklingError,
-                    ) as exc:
-                        logger.debug(
-                            "Cannot read broker message %s",
-                            exc,
-                            exc_info=True,
-                        )
+                if not message:
+                    continue
+                yield message["data"]
+
+
+class ListQueueBroker(BaseRedisBroker):
+    """Broker that works with Redis and distributes tasks between workers."""
+
+    async def kick(self, message: BrokerMessage) -> None:  # noqa: D102
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            await redis_conn.lpush(self.queue_name, pickle.dumps(message))
+
+    async def _listen_to_raw_messages(self) -> AsyncGenerator[bytes, None]:
+        redis_brpop_data_position = 1
+        async with Redis(connection_pool=self.connection_pool) as redis_conn:
+            while True:  # noqa: WPS457
+                yield (await redis_conn.brpop(self.queue_name))[
+                    redis_brpop_data_position
+                ]

taskiq_redis-0.1.0/taskiq_redis/__init__.py

@@ -1,5 +0,0 @@
-"""Package for redis integration."""
-from taskiq_redis.redis_backend import RedisAsyncResultBackend
-from taskiq_redis.redis_broker import RedisBroker
-
-__all__ = ["RedisAsyncResultBackend", "RedisBroker"]

taskiq_redis-0.1.0/taskiq_redis/redis_backend.py

@@ -1,98 +0,0 @@
-import pickle
-from typing import TypeVar
-
-from redis.asyncio import ConnectionPool, Redis
-from taskiq import AsyncResultBackend
-from taskiq.abc.result_backend import TaskiqResult
-
-_ReturnType = TypeVar("_ReturnType")
-
-
-class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
-    """Async result based on redis."""
-
-    def __init__(self, redis_url: str, keep_results: bool = True):
-        """
-        Constructs a new result backend.
-
-        :param redis_url: url to redis.
-        :param keep_results: flag to not remove results from Redis after reading.
-        """
-        self.redis_pool = ConnectionPool.from_url(redis_url)
-        self.keep_results = keep_results
-
-    async def shutdown(self) -> None:
-        """Closes redis connection."""
-        await self.redis_pool.disconnect()
-
-    async def set_result(
-        self,
-        task_id: str,
-        result: TaskiqResult[_ReturnType],
-    ) -> None:
-        """
-        Sets task result in redis.
-
-        Dumps TaskiqResult instance into the bytes and writes
-        it to redis.
-
-        :param task_id: ID of the task.
-        :param result: TaskiqResult instance.
-        """
-        result_dict = result.dict(exclude={"return_value"})
-
-        for result_key, result_value in result_dict.items():
-            result_dict[result_key] = pickle.dumps(result_value)
-        # This trick will preserve original returned value.
-        # It helps when you return not serializable classes.
-        result_dict["return_value"] = pickle.dumps(result.return_value)
-
-        async with Redis(connection_pool=self.redis_pool) as redis:
-            await redis.hset(
-                task_id,
-                mapping=result_dict,
-            )
-
-    async def is_result_ready(self, task_id: str) -> bool:
-        """
-        Returns whether the result is ready.
-
-        :param task_id: ID of the task.
-
-        :returns: True if the result is ready else False.
-        """
-        async with Redis(connection_pool=self.redis_pool) as redis:
-            return bool(await redis.exists(task_id))
-
-    async def get_result(  # noqa: WPS210
-        self,
-        task_id: str,
-        with_logs: bool = False,
-    ) -> TaskiqResult[_ReturnType]:
-        """
-        Gets result from the task.
-
-        :param task_id: task's id.
-        :param with_logs: if True it will download task's logs.
-        :return: task's return value.
-        """
-        fields = list(TaskiqResult.__fields__.keys())
-
-        if not with_logs:
-            fields.remove("log")
-
-        async with Redis(connection_pool=self.redis_pool) as redis:
-            result_values = await redis.hmget(
-                name=task_id,
-                keys=fields,
-            )
-
-            if not self.keep_results:
-                await redis.delete(task_id)
-
-        result = {
-            result_key: pickle.loads(result_value)
-            for result_value, result_key in zip(result_values, fields)
-        }
-
-        return TaskiqResult(**result)