feat: add new result_backend with capo-s3 (#40)

Author: danfimov

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "taskiq>=0.12.6",
     "aiobotocore>=2.13.3",
     "capo-s3>=0.15.0",
+    "capo-sqs>=0.6.0",
 ]
 
 [project.urls]
@@ -43,7 +44,8 @@ dependencies = [
 dev = [
     {include-group = "lint"},
     {include-group = "test"},
-    {include-group = "types"},
+    {include-group = "examples"},
+    {include-group = "docs"},
     "prek>=0.5.3",
 ]
 test = [
@@ -54,14 +56,15 @@ test = [
 lint = [
     "ruff>=0.16.7",
     "zizmor>=1.30.1",
-]
-types = [
     "mypy>=2.3.1",
     "types-aiobotocore[essential]>=3.7.0",
 ]
 examples = [
     "python-dotenv>=1.2.3",
 ]
+docs = [
+    "zensical>=0.0.62",
+]
 
 
 [build-system]

src/taskiq_sqs/result_backend.py

@@ -1,7 +1,7 @@
-from typing import TYPE_CHECKING, Any, TypeVar
+import contextlib
+from typing import Any, TypeVar
 
-from aiobotocore.session import get_session
-from botocore.exceptions import ClientError
+import capo_s3
 from taskiq import AsyncResultBackend
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
@@ -12,9 +12,6 @@
 from taskiq_sqs.types import S3Bucket
 
 
-if TYPE_CHECKING:
-    from types_aiobotocore_s3.client import S3Client
-
 _ReturnType = TypeVar("_ReturnType")
 
 
@@ -48,61 +45,57 @@ def __init__(
         self._aws_secret_access_key = aws_secret_access_key
         self._bucket = bucket
         self._base_path = base_path
-        self._session = get_session()
         self._serializer = serializer or JSONSerializer()
 
-    async def _get_client(self) -> "S3Client":
-        """
-        Retrieves the S3 client, creating it if necessary.
-
-        Returns:
-            S3Client: The initialized S3 client.
-        """
-        self._client_context_creator = self._session.create_client(
-            "s3",
-            region_name=self._aws_region,
-            endpoint_url=self._aws_endpoint_url,
-            aws_access_key_id=self._aws_access_key_id,
-            aws_secret_access_key=self._aws_secret_access_key,
-        )
-        return await self._client_context_creator.__aenter__()
-
     async def startup(self) -> None:
-        """Initialize the result backend."""
-        self._s3_client = await self._get_client()
+        """Initialize the S3 client and ensure the bucket exists."""
+        credentials = None
+        if self._aws_access_key_id and self._aws_secret_access_key:
+            credentials = capo_s3.Credentials(
+                access_key=self._aws_access_key_id,
+                secret_key=self._aws_secret_access_key,
+            )
+        self._s3_client = capo_s3.AsyncS3Client(
+            region=self._aws_region,
+            endpoint=self._aws_endpoint_url,
+            credentials=credentials,
+            force_path_style=True,
+        )
+        await self._s3_client.__aenter__()
         try:
             await self._ensure_bucket_exists()
         except Exception:
-            await self._client_context_creator.__aexit__(None, None, None)
+            await self._s3_client.__aexit__(None, None, None)
             raise
         return await super().startup()
 
     async def _ensure_bucket_exists(self) -> None:
         try:
-            await self._s3_client.head_bucket(Bucket=self._bucket["name"])
-        except ClientError as exc:
-            code = exc.response.get("Error", {}).get("Code")
-            if code not in ("404", "NoSuchBucket"):
-                raise exceptions.ResultBackendError(code=code) from exc
+            await self._s3_client.head_bucket(bucket=self._bucket["name"])
+        except capo_s3.errors.NotFound:
             if not self._bucket.get("declare", True):
-                raise exceptions.BucketNotFoundError(bucket_name=self._bucket["name"]) from exc
+                raise exceptions.BucketNotFoundError(bucket_name=self._bucket["name"]) from None
             await self._create_bucket()
+        except capo_s3.errors.ServiceError as exc:
+            raise exceptions.ResultBackendError(code=exc.code) from exc
 
     async def _create_bucket(self) -> None:
-        create_kwargs: dict[str, Any] = {"Bucket": self._bucket["name"]}
+        create_kwargs: dict[str, Any] = {}
         if self._aws_region and self._aws_region != constants.AWS_DEFAULT_REGION:
-            create_kwargs["CreateBucketConfiguration"] = {"LocationConstraint": self._aws_region}
-        try:
-            await self._s3_client.create_bucket(**create_kwargs)
-        except ClientError as exc:
-            if exc.response.get("Error", {}).get("Code") != "BucketAlreadyOwnedByYou":  # can be raise between workers
-                raise
+            create_kwargs["create_bucket_configuration"] = {"location_constraint": self._aws_region}
+        with contextlib.suppress(capo_s3.errors.BucketAlreadyOwnedByYou):
+            await self._s3_client.create_bucket(bucket=self._bucket["name"], **create_kwargs)
 
     async def shutdown(self) -> None:
         """Shut down the result backend."""
-        await self._client_context_creator.__aexit__(None, None, None)
+        await self._s3_client.__aexit__(None, None, None)
         return await super().shutdown()
 
+    def _build_key(self, task_id: str) -> str:
+        if self._base_path:
+            return f"{self._base_path.rstrip('/')}/{task_id}"
+        return task_id
+
     async def set_result(
         self,
         task_id: str,
@@ -114,13 +107,10 @@ async def set_result(
         :param task_id: current task id.
         :param result: result of execution.
         """
-        if self._base_path:
-            task_id = f"{self._base_path.rstrip('/')}/{task_id}"
-
         await self._s3_client.put_object(
-            Bucket=self._bucket["name"],
-            Key=task_id,
-            Body=self._serializer.dumpb(model_dump(result)),
+            bucket=self._bucket["name"],
+            key=self._build_key(task_id),
+            body=self._serializer.dumpb(model_dump(result)),
         )
 
     async def get_result(
@@ -138,27 +128,17 @@ async def get_result(
         :param with_logs: whether to fetch logs.
         :return: result.
         """
-        result = None
-        if self._base_path:
-            task_id = f"{self._base_path.rstrip('/')}/{task_id}"
         try:
-            if response := await self._s3_client.get_object(
-                Bucket=self._bucket["name"],
-                Key=task_id,
-            ):
-                async with response["Body"] as stream:
-                    result = await stream.read()
-        except ClientError as exc:
-            code = exc.response.get("Error", {}).get("Code")
-            if code in ["NoSuchKey", "404"]:
-                raise exceptions.ResultIsMissingError(task_id=task_id) from exc
-            raise exceptions.ResultBackendError(code=code) from exc
-        if result is None:
-            raise exceptions.ResultIsMissingError(task_id=task_id)
+            async with self._s3_client.get_object(bucket=self._bucket["name"], key=self._build_key(task_id)) as output:
+                body = b"".join([chunk async for chunk in output["body"]])
+        except capo_s3.errors.NoSuchKey as exc:
+            raise exceptions.ResultIsMissingError(task_id=task_id) from exc
+        except capo_s3.errors.ServiceError as exc:
+            raise exceptions.ResultBackendError(code=exc.code) from exc
 
         taskiq_result = model_validate(
             TaskiqResult[_ReturnType],
-            self._serializer.loadb(result),
+            self._serializer.loadb(body),
         )
 
         if not with_logs:
@@ -173,17 +153,10 @@ async def is_result_ready(self, task_id: str) -> bool:
         :param task_id: id of a task.
         :return: True if result is ready.
         """
-        if self._base_path:
-            task_id = f"{self._base_path.rstrip('/')}/{task_id}"
         try:
-            if await self._s3_client.head_object(Bucket=self._bucket["name"], Key=task_id):
-                return True
-        except ClientError as exc:
-            code = exc.response.get("Error", {}).get("Code")
-            if code in ["NoSuchKey", "404"]:
-                pass
-            else:
-                raise exceptions.ResultBackendError(
-                    code=code,
-                ) from exc
-        return False
+            await self._s3_client.head_object(bucket=self._bucket["name"], key=self._build_key(task_id))
+        except capo_s3.errors.NotFound:
+            return False
+        except capo_s3.errors.ServiceError as exc:
+            raise exceptions.ResultBackendError(code=exc.code) from exc
+        return True

tests/test_result_backend.py

@@ -1,7 +1,6 @@
 import uuid
 from collections.abc import AsyncGenerator
 from typing import TYPE_CHECKING, Any
-from unittest.mock import AsyncMock
 
 import pytest
 from taskiq.result import TaskiqResult
@@ -26,12 +25,13 @@ class TestResultBackend:
     async def test_when_result_set__then_result_is_actually_saved_to_s3(
         self,
         s3_backend: S3ResultBackend,
+        s3_client: "S3Client",
         s3_bucket: str,
         taskiq_result: TaskiqResult,
     ) -> None:
         await s3_backend.set_result("test_task_id", taskiq_result)
 
-        response = await s3_backend._s3_client.get_object(
+        response = await s3_client.get_object(
             Bucket=s3_bucket,
             Key="test_task_id",
         )
@@ -48,24 +48,17 @@ async def test_when_result_present_in_s3__then_get_result_return_it(
         assert retrieved_result.return_value == "test_value"
         assert retrieved_result.is_err is True
 
-    async def test_when_result_is_missing__then_get_result_raise_exception(
-        self,
-        s3_backend: S3ResultBackend,
-    ) -> None:
-        s3_backend._s3_client.get_object = AsyncMock(return_value={})  # Simulate a response with no Body
-        with pytest.raises(ResultIsMissingError):
-            await s3_backend.get_result("test_task_id")
-
     async def test_when_set_result_is_called__then_save_it_to_right_path(
         self,
         s3_backend: S3ResultBackend,
+        s3_client: "S3Client",
         s3_bucket: str,
         taskiq_result: TaskiqResult,
     ) -> None:
         s3_backend._base_path = "results"
         await s3_backend.set_result("test_task_id", taskiq_result)
 
-        response = await s3_backend._s3_client.head_object(
+        response = await s3_client.head_object(
             Bucket=s3_bucket,
             Key="results/test_task_id",
         )