Merge pull request #10 from ApeWorX/feat/ecs-creds

feat: adds force_ecs_container_credentials kwarg to broker

Author: Mike Shultz

taskiq_sqs/aws.py

@@ -0,0 +1,22 @@
+import json
+import os
+
+import urllib3  # boto3 peer dep (v1)
+
+ECS_CONTAINER_METADATA_URI = "http://169.254.170.2"
+
+
+class InvalidEnvironment(Exception):
+    pass
+
+
+def get_container_credentials():
+    """Fetches the ECS task role credentials provided by the metadata service"""
+    if not (relative_uri := os.environ.get("AWS_CONTAINER_CREDENTIALS_RELATIVE_URI")):
+        raise InvalidEnvironment(
+            "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI not defined. This may not be an ECS container."
+        )
+
+    http = urllib3.PoolManager()
+    resp = http.request("GET", f"{ECS_CONTAINER_METADATA_URI}{relative_uri}")
+    return json.loads(resp.data)

taskiq_sqs/broker.py

@@ -16,6 +16,8 @@
 from taskiq.exceptions import BrokerError
 from taskiq.message import BrokerMessage
 
+from taskiq_sqs.aws import get_container_credentials
+
 if TYPE_CHECKING:
     from mypy_boto3_sqs.service_resource import Queue, SQSServiceResource
 
@@ -37,15 +39,28 @@ def __init__(
         result_backend: Optional[AsyncResultBackend] = None,
         task_id_generator: Optional[Callable[[], str]] = None,
         sqs_region_override: str | None = None,
+        force_ecs_container_credentials=False,
     ) -> None:
         super().__init__(result_backend, task_id_generator)
 
         if not sqs_queue_url or not sqs_queue_url.startswith("http"):
             raise BrokerError("A valid SQS Queue URL is required")
 
+        creds = dict()
+        # NOTE: This bypasses the normal order of operations for boto3 auth and
+        #       goes straight to using the ECS role creds from the metadata
+        #       service. This can be useful in edge cases where there are higher
+        #       priority credentials you do not want to use for this service.
+        if force_ecs_container_credentials:
+            creds = get_container_credentials()
+
         self.sqs_queue_url = sqs_queue_url
         self._sqs: SQSServiceResource = boto3.resource(
-            "sqs", region_name=sqs_region_override
+            "sqs",
+            region_name=sqs_region_override,
+            aws_access_key_id=creds.get("AccessKeyId"),
+            aws_secret_access_key=creds.get("SecretAccessKey"),
+            aws_session_token=creds.get("Token"),
         )
         self._sqs_queue: Optional[Queue] = None