Merge pull request #1 from ApeWorX/refactor/name

fix: name and other stuff

Author: Mike Shultz

.github/workflows/test.yaml

@@ -9,6 +9,9 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
   cancel-in-progress: true
 
+env:
+  AWS_DEFAULT_REGION: "us-east-1"
+
 jobs:
     linting:
         runs-on: ubuntu-latest
@@ -61,7 +64,7 @@ jobs:
         strategy:
             matrix:
                 os: [ubuntu-latest, macos-latest]   # eventually add `windows-latest`
-                python-version: ["3.8", "3.9", "3.10", "3.11"]
+                python-version: ["3.10", "3.11", "3.12"]
 
         steps:
         - uses: actions/checkout@v3

README.md

@@ -5,3 +5,32 @@ Mostly generic SQS async broker for TaskIQ.
 ## Expiration
 
 If you set the `sqs_expiry` label to a unix timestamp, the message will be discarded if the worker receives it after that time.
+
+```python
+import asyncio
+from taskiq_sqs import SQSBroker
+
+broker = SQSBroker("http://sqs.us-east-1.localhost.localstack.cloud:4566/000000000000/my-queue")
+
+@broker.task
+async def add_one(value: int) -> int:
+    return value + 1
+
+
+async def main() -> None:
+    # Never forget to call startup in the beginning.
+    await broker.startup()
+    # Send the task to the broker.
+    task = await add_one.kiq(1)
+    # Wait for the result. (result backend must be configured)
+    result = await task.wait_result(timeout=2)
+    print(f"Task execution took: {result.execution_time} seconds.")
+    if not result.is_err:
+        print(f"Returned value: {result.return_value}")
+    else:
+        print("Error found while executing task.")
+    await broker.shutdown()
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```

pyproject.toml

@@ -7,7 +7,7 @@ name = "taskiq-sqs"
 description = "SQS Broker for TaskIQ"
 urls = {source = "https://github.com/ApeWorx/taskiq-sqs"}
 readme = "README.md"
-requires-python = ">=3.8"
+requires-python = ">=3.10"
 keywords = ["taskiq", "broker", "aws", "sqs"]
 license = {file = "LICENSE"}
 classifiers = [
@@ -41,6 +41,7 @@ dev = [
     "pytest~=7.4.4",
     "requests~=2.31.0",
     "setuptools-scm~=8.0.4",
+    "taskiq>=0.10.3,<1",
     "types-requests>=2.31.0.20240106",
 ]
 

taskiq_redis/__init__.py

@@ -0,0 +1,3 @@
+from taskiq_sqs.broker import SQSBroker
+
+__all__ = ["SQSBroker"]

taskiq_sqs/__init__.py

@@ -1,15 +1,9 @@
 import asyncio
 import logging
 from datetime import datetime, timezone
-from typing import (
-    AsyncGenerator,
-    Callable,
-    Optional,
-    Union,
-)
+from typing import AsyncGenerator, Callable, Optional, Union
 
 import boto3
-from ape.logging import logger
 from asyncer import asyncify
 from mypy_boto3_sqs.service_resource import Queue
 from taskiq import AsyncBroker
@@ -18,7 +12,10 @@
 from taskiq.message import BrokerMessage
 
 logger = logging.getLogger(__name__)
-stamp = lambda: int(datetime.now(tz=timezone.utc).timestamp())
+
+
+def stamp() -> int:
+    return int(datetime.now(tz=timezone.utc).timestamp())
 
 
 class SQSBroker(AsyncBroker):
@@ -39,7 +36,9 @@ async def _get_queue(self) -> Queue:
         queue_name = self.sqs_queue_url.split("/")[-1]
 
         if not self._sqs_queue:
-            self._sqs_queue = await asyncify(self._sqs.get_queue_by_name)(QueueName=queue_name)
+            self._sqs_queue = await asyncify(self._sqs.get_queue_by_name)(
+                QueueName=queue_name
+            )
 
             if not self._sqs_queue:
                 raise Exception("SQS Queue not found")
@@ -63,18 +62,16 @@ async def kick(
         """
         queue = await self._get_queue()
         # Must be explicitly set as a label to a unix timestamp
-        expiry = message.labels.pop("sqs_expiry", None)
-
-        # SQS structured message attributes
-        message_attributes = dict()
-        if expiry:
-            message_attributes["expiry"] = {
-                "StringValue": str(expiry),
-                "DataType": "Number",
-            }
+        expiry = message.labels.pop("sqs_expiry", 0)
 
         await asyncify(queue.send_message)(
-            MessageAttributes=message_attributes,
+            # SQS structured message attributes
+            MessageAttributes={
+                "expiry": {
+                    "StringValue": str(expiry),
+                    "DataType": "Number",
+                }
+            },
             MessageBody=message.message.decode("utf-8"),
             MessageGroupId=message.task_name,
         )
@@ -102,14 +99,20 @@ async def listen(self) -> AsyncGenerator[Union[bytes, AckableMessage], None]:
         while True:
             last_had_message = False
 
-            for message in await asyncify(queue.receive_messages)():
+            for message in await asyncify(queue.receive_messages)(
+                MessageAttributeNames=[".*"]
+            ):
                 try:
                     if message.message_attributes:
                         # if expiry was set as a message attribute, respect it
-                        if expiry := message.message_attributes.get("expiry"):
-                            diff = stamp() - expiry
-                            if diff > 0:
-                                logger.debug(f"Message expired {diff} seconds ago. Skipping.")
+                        if expiry_typed := message.message_attributes.get("expiry"):
+                            expiry = int(expiry_typed.get("StringValue", 0))
+                            now = stamp()
+                            if 0 < expiry < now:
+                                logger.warn(
+                                    f"Message expired {now - expiry} seconds ago. Skipping."
+                                )
+                                await asyncify(message.delete)()
                                 continue
                 except TypeError:
                     # Ignore weird expiries.  Not critical.

taskiq_redis/broker.py taskiq_sqs/broker.pytaskiq_redis/broker.py renamed to taskiq_sqs/broker.py

@@ -0,0 +1,11 @@
+from taskiq_sqs import SQSBroker
+
+
+def test_init():
+    broker = SQSBroker("https://sqs.us-west-2.amazonaws.com/123456789012/queue-name")
+    assert (
+        broker.sqs_queue_url
+        == "https://sqs.us-west-2.amazonaws.com/123456789012/queue-name"
+    )
+    assert broker._sqs is not None
+    assert broker._sqs_queue is None