feat: add S3 result backend

Author: danfimov

examples/example_broker.py

@@ -10,9 +10,8 @@
 
 import boto3
 import dotenv
-from taskiq_redis import RedisAsyncResultBackend
 
-from taskiq_sqs import SQSBroker
+from taskiq_sqs import S3Bucket, S3ResultBackend, SQSBroker
 
 
 dotenv.load_dotenv()
@@ -24,7 +23,7 @@
 boto3.client("sqs").create_queue(QueueName=QUEUE_NAME)
 
 broker = SQSBroker(QUEUE_URL, sqs_region_override="us-east-1").with_result_backend(
-    RedisAsyncResultBackend(redis_url="redis://localhost:6379")
+    S3ResultBackend(bucket=S3Bucket(name="response-bucket"))
 )
 
 

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
     "taskiq>=0.12.1",
     "asyncer~=0.0.5",
     "boto3~=1.34.34",
+    "aiobotocore>=2.13.3",
 ]
 
 # [project.urls]
@@ -60,6 +61,7 @@ types = [
     "mypy>=2.1.0",
     "mypy-boto3-sqs>=1.34.101",
     "boto3-stubs[essential]>=1.34.84",
+    "types-aiobotocore[essential]>=3.7.0",
 ]
 examples = [
     "python-dotenv>=1.2.2",
@@ -139,6 +141,7 @@ ignore = [
     "D",
     "S101",  # assert usage
     "SLF001",  # private member accessed
+    "PLR2004",  # magic values
 ]
 "examples/*" = [
     "T201",  # print

src/taskiq_sqs/__init__.py

@@ -1,4 +1,10 @@
 from taskiq_sqs.broker import SQSBroker
+from taskiq_sqs.bucket import S3Bucket
+from taskiq_sqs.result_backend import S3ResultBackend
 
 
-__all__ = ["SQSBroker"]
+__all__ = [
+    "S3Bucket",
+    "S3ResultBackend",
+    "SQSBroker",
+]

src/taskiq_sqs/broker.py

@@ -14,6 +14,7 @@
 from taskiq.message import BrokerMessage
 
 from taskiq_sqs.aws import get_container_credentials
+from taskiq_sqs.exceptions import BrokerInitError
 
 
 if TYPE_CHECKING:
@@ -27,10 +28,6 @@ def stamp() -> int:  # noqa: D103
     return int(datetime.now(tz=timezone.utc).timestamp())
 
 
-class SQSBrokerError(Exception):
-    """Generic SQS broker error."""
-
-
 class SQSBroker(AsyncBroker):
     """AWS SQS TaskIQ broker."""
 
@@ -47,7 +44,7 @@ def __init__(  # noqa: D107
         super().__init__(result_backend, task_id_generator)
 
         if not sqs_queue_url or not sqs_queue_url.startswith("http"):
-            raise SQSBrokerError("A valid SQS queue url is required")
+            raise BrokerInitError(details="A valid SQS queue url is required")
 
         # NOTE: This bypasses the normal order of operations for boto3 auth and
         #       goes straight to using the ECS role creds from the metadata
@@ -61,7 +58,7 @@ def __init__(  # noqa: D107
         self._creds_expiration: datetime | None = None
 
         if max_number_of_messages > 10:  # noqa: PLR2004
-            raise SQSBrokerError("MaxNumberOfMessages can be no greater than 10")
+            raise BrokerInitError(details="MaxNumberOfMessages can be no greater than 10")
 
         self.wait_time_seconds = max(wait_time_seconds, 0)
         self.max_number_of_messages = max(max_number_of_messages, 1)
@@ -100,8 +97,7 @@ async def _get_queue(self) -> "Queue":
         )
 
         if not self._sqs_queue:
-            exc_message = "SQS queue not found"
-            raise Exception(exc_message)  # noqa: TRY002
+            raise BrokerInitError(details="SQS queue not found")
 
         return self._sqs_queue
 

src/taskiq_sqs/bucket.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class S3Bucket:
+    """
+    Represents an S3 bucket configuration.
+
+    Attributes:
+        name: The name of the bucjet.
+        declare: Whether to create the bucket on startup if it not exists yet.
+    """
+    name: str
+    declare: bool = True

src/taskiq_sqs/constants.py

@@ -0,0 +1 @@
+AWS_DEFAULT_REGION = "us-east-1"

src/taskiq_sqs/exceptions.py

@@ -0,0 +1,32 @@
+from taskiq.exceptions import TaskiqError
+
+
+class BaseTaskiqSQSError(TaskiqError):
+    """Base error from taskiq-sqs."""
+
+
+class BrokerInitError(BaseTaskiqSQSError):
+    """Error during broker initialization."""
+
+    __template__ = "Error during broker initialization: {details}"
+    details: str
+
+
+class ResultBackendError(BaseTaskiqSQSError):
+    """Base error for all taskiq-aio-sqs broker exceptions."""
+
+    __template__ = "Unexpected error occurred: {code}"
+    code: str | None = None
+
+
+class BucketNotFoundError(BaseTaskiqSQSError):
+    """Error if bucket not found."""
+
+    __template__ = "Bucket '{bucket_name}' not found during initialization and declare=False"
+    bucket_name: str
+
+class ResultIsMissingError(BaseTaskiqSQSError):
+    """Error if there is no result when we trying to get it."""
+
+    __template__ = "Result for task {task_id} is missing in the result backend"
+    task_id: str

src/taskiq_sqs/result_backend.py

@@ -0,0 +1,186 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from aiobotocore.session import get_session
+from botocore.exceptions import ClientError
+from taskiq import AsyncResultBackend
+from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.compat import model_dump, model_validate
+from taskiq.result import TaskiqResult
+from taskiq.serializers import JSONSerializer
+
+from taskiq_sqs import constants, exceptions
+from taskiq_sqs.bucket import S3Bucket
+
+
+if TYPE_CHECKING:
+    from types_aiobotocore_s3.client import S3Client
+
+_ReturnType = TypeVar("_ReturnType")
+
+
+class S3ResultBackend(AsyncResultBackend[_ReturnType]):
+    """TaskIQ result backend that uses S3."""
+
+    def __init__(
+        self,
+        bucket: S3Bucket,
+        base_path: str = "",
+        endpoint_url: str | None = None,
+        aws_region_name: str = constants.AWS_DEFAULT_REGION,
+        aws_access_key_id: str | None = None,
+        aws_secret_access_key: str | None = None,
+        serializer: TaskiqSerializer | None = None,
+    ) -> None:
+        """
+        Constructs a new S3 result backend.
+
+        :param bucket_name: name of the bucket on S3.
+        :param base_path: base path for results.
+        :param endpoint_url: endpoint URL for S3.
+        :param aws_region_name: AWS region, default is 'us-east-1'.
+        :param aws_access_key_id: AWS access key ID.
+        :param aws_secret_access_key: AWS secret access key.
+        :param serializer: serializer to use.
+        """
+        self._aws_region = aws_region_name
+        self._aws_endpoint_url = endpoint_url
+        self._aws_access_key_id = aws_access_key_id
+        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()
+        await self._ensure_bucket_exists()
+        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
+            if not self._bucket.declare:
+                raise exceptions.BucketNotFoundError(bucket_name=self._bucket.name) from exc
+            await self._create_bucket()
+
+    async def _create_bucket(self) -> None:
+        create_kwargs: dict[str, Any] = {"Bucket": self._bucket.name}
+        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
+
+    async def shutdown(self) -> None:
+        """Shut down the result backend."""
+        await self._client_context_creator.__aexit__(None, None, None)
+        return await super().shutdown()
+
+    async def set_result(
+        self,
+        task_id: str,
+        result: TaskiqResult[_ReturnType],
+    ) -> None:
+        """
+        Set result in your backend.
+
+        :param task_id: current task id.
+        :param result: result of execution.
+        """
+        if self._base_path:
+            task_id = str(Path(self._base_path) / task_id)
+
+        await self._s3_client.put_object(
+            Bucket=self._bucket.name,
+            Key=task_id,
+            Body=self._serializer.dumpb(model_dump(result)),
+        )
+
+    async def get_result(
+        self,
+        task_id: str,
+        with_logs: bool = False,
+    ) -> TaskiqResult[_ReturnType]:
+        """
+        Here you must retrieve result by id.
+
+        Logs is a part of a result. Here we have a parameter whether you want to fetch result with logs or not,
+        because logs can have a lot of info and sometimes it's critical to get only needed information.
+
+        :param task_id: id of a task.
+        :param with_logs: whether to fetch logs.
+        :return: result.
+        """
+        result = None
+        if self._base_path:
+            task_id = str(Path(self._base_path) / 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)
+
+        taskiq_result = model_validate(
+            TaskiqResult[_ReturnType],
+            self._serializer.loadb(result),
+        )
+
+        if not with_logs:
+            taskiq_result.log = None
+
+        return taskiq_result
+
+    async def is_result_ready(self, task_id: str) -> bool:
+        """
+        Check if result exists.
+
+        :param task_id: id of a task.
+        :return: True if result is ready.
+        """
+        if self._base_path:
+            task_id = str(Path(self._base_path) / 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