taskiq-redis 1.0.0__tar.gz → 1.0.1__tar.gz

{taskiq_redis-1.0.0 → taskiq_redis-1.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: taskiq-redis
-Version: 1.0.0
+Version: 1.0.1
 Summary: Redis integration for taskiq
 Home-page: https://github.com/taskiq-python/taskiq-redis
 Keywords: taskiq,tasks,distributed,async,redis,result_backend

{taskiq_redis-1.0.0 → taskiq_redis-1.0.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-redis"
-version = "1.0.0"
+version = "1.0.1"
 description = "Redis integration for taskiq"
 authors = ["taskiq-team <taskiq@norely.com>"]
 readme = "README.md"
@@ -39,7 +39,7 @@ pytest-env = "^0.6.2"
 fakeredis = "^2"
 pre-commit = "^2.20.0"
 pytest-xdist = { version = "^2.5.0", extras = ["psutil"] }
-ruff = "^0.1.0"
+ruff = "^0"
 types-redis = "^4.6.0.20240425"
 
 [tool.mypy]

{taskiq_redis-1.0.0 → taskiq_redis-1.0.1}/taskiq_redis/redis_backend.py

@@ -19,6 +19,7 @@ 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.depends.progress_tracker import TaskProgress
 from taskiq.serializers import PickleSerializer
 
 from taskiq_redis.exceptions import (
@@ -41,6 +42,8 @@ else:
 
 _ReturnType = TypeVar("_ReturnType")
 
+PROGRESS_KEY_SUFFIX = "__progress"
+
 
 class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
     """Async result based on redis."""
@@ -174,6 +177,55 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
 
         return taskiq_result
 
+    async def set_progress(
+        self,
+        task_id: str,
+        progress: TaskProgress[_ReturnType],
+    ) -> None:
+        """
+        Sets task progress in redis.
+
+        Dumps TaskProgress instance into the bytes and writes
+        it to redis with a standard suffix on the task_id as the key
+
+        :param task_id: ID of the task.
+        :param result: task's TaskProgress instance.
+        """
+        redis_set_params: Dict[str, Union[str, int, bytes]] = {
+            "name": task_id + PROGRESS_KEY_SUFFIX,
+            "value": self.serializer.dumpb(model_dump(progress)),
+        }
+        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)  # type: ignore
+
+    async def get_progress(
+        self,
+        task_id: str,
+    ) -> Union[TaskProgress[_ReturnType], None]:
+        """
+        Gets progress results from the task.
+
+        :param task_id: task's id.
+        :return: task's TaskProgress instance.
+        """
+        async with Redis(connection_pool=self.redis_pool) as redis:
+            result_value = await redis.get(
+                name=task_id + PROGRESS_KEY_SUFFIX,
+            )
+
+        if result_value is None:
+            return None
+
+        return model_validate(
+            TaskProgress[_ReturnType],
+            self.serializer.loadb(result_value),
+        )
+
 
 class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
     """Async result backend based on redis cluster."""
@@ -301,6 +353,53 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
 
         return taskiq_result
 
+    async def set_progress(
+        self,
+        task_id: str,
+        progress: TaskProgress[_ReturnType],
+    ) -> None:
+        """
+        Sets task progress in redis.
+
+        Dumps TaskProgress instance into the bytes and writes
+        it to redis with a standard suffix on the task_id as the key
+
+        :param task_id: ID of the task.
+        :param result: task's TaskProgress instance.
+        """
+        redis_set_params: Dict[str, Union[str, int, bytes]] = {
+            "name": task_id + PROGRESS_KEY_SUFFIX,
+            "value": self.serializer.dumpb(model_dump(progress)),
+        }
+        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
+
+        await self.redis.set(**redis_set_params)  # type: ignore
+
+    async def get_progress(
+        self,
+        task_id: str,
+    ) -> Union[TaskProgress[_ReturnType], None]:
+        """
+        Gets progress results from the task.
+
+        :param task_id: task's id.
+        :return: task's TaskProgress instance.
+        """
+        result_value = await self.redis.get(  # type: ignore[attr-defined]
+            name=task_id + PROGRESS_KEY_SUFFIX,
+        )
+
+        if result_value is None:
+            return None
+
+        return model_validate(
+            TaskProgress[_ReturnType],
+            self.serializer.loadb(result_value),
+        )
+
 
 class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
     """Async result based on redis sentinel."""
@@ -439,6 +538,55 @@ class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
 
         return taskiq_result
 
+    async def set_progress(
+        self,
+        task_id: str,
+        progress: TaskProgress[_ReturnType],
+    ) -> None:
+        """
+        Sets task progress in redis.
+
+        Dumps TaskProgress instance into the bytes and writes
+        it to redis with a standard suffix on the task_id as the key
+
+        :param task_id: ID of the task.
+        :param result: task's TaskProgress instance.
+        """
+        redis_set_params: Dict[str, Union[str, int, bytes]] = {
+            "name": task_id + PROGRESS_KEY_SUFFIX,
+            "value": self.serializer.dumpb(model_dump(progress)),
+        }
+        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 self._acquire_master_conn() as redis:
+            await redis.set(**redis_set_params)  # type: ignore
+
+    async def get_progress(
+        self,
+        task_id: str,
+    ) -> Union[TaskProgress[_ReturnType], None]:
+        """
+        Gets progress results from the task.
+
+        :param task_id: task's id.
+        :return: task's TaskProgress instance.
+        """
+        async with self._acquire_master_conn() as redis:
+            result_value = await redis.get(
+                name=task_id + PROGRESS_KEY_SUFFIX,
+            )
+
+        if result_value is None:
+            return None
+
+        return model_validate(
+            TaskProgress[_ReturnType],
+            self.serializer.loadb(result_value),
+        )
+
     async def shutdown(self) -> None:
         """Shutdown sentinel connections."""
         for sentinel in self.sentinel.sentinels: