Merge branch 'taskiq-python:main' into anton/try_except_listen

Author: khizunov

.flake8

@@ -18,6 +18,8 @@ jobs:
         uses: knowsuchagency/poetry-install@v1
         env:
           POETRY_VIRTUALENVS_CREATE: false
+      - name: Set version
+        run: poetry version "${{ github.ref_name }}"
       - name: Release package
         env:
           POETRY_PYPI_TOKEN_PYPI: ${{ secrets.PYPI_TOKEN }}

.github/workflows/release.yml

@@ -37,7 +37,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Set up Redis instance and Redis cluster
-        run: docker-compose up -d
+        run: docker compose up -d
       - name: Set up Python
         uses: actions/setup-python@v2
         with:

.github/workflows/test.yml

@@ -18,7 +18,7 @@ repos:
     hooks:
       - id: black
         name: Format with Black
-        entry: black
+        entry: poetry run black
         language: system
         types: [python]
 
@@ -29,12 +29,13 @@ repos:
         pass_filenames: false
         types: [python]
         args:
+          - check
           - "--fix"
           - "taskiq_redis"
           - "tests"
 
       - id: mypy
         name: Validate types with MyPy
-        entry: mypy
+        entry: poetry run mypy
         language: system
         types: [ python ]

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-redis"
-version = "1.0.0"
+version = "0.0.0"
 description = "Redis integration for taskiq"
 authors = ["taskiq-team <[email protected]>"]
 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]

poetry.lock

@@ -19,6 +19,7 @@
 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 @@
 
 _ReturnType = TypeVar("_ReturnType")
 
+PROGRESS_KEY_SUFFIX = "__progress"
+
 
 class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
     """Async result based on redis."""
@@ -174,6 +177,55 @@ async def get_result(
 
         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 @@ async def get_result(
 
         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 @@ async def get_result(
 
         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: