remember to get aws credentials first
- ID
f8cf567- date
2023-06-21 23:45:08+00:00- author
Alex Chan <alex@alexwlchan.net>- parent
2890f1b- message
remember to get aws credentials first- changed files
2 files, 89 additions, 78 deletions
Changed files
aws/download_sqs_messages (2065) → aws/download_sqs_messages (116)
diff --git a/aws/download_sqs_messages b/aws/download_sqs_messages
index 46e19a4..765fd8e 100755
--- a/aws/download_sqs_messages
+++ b/aws/download_sqs_messages
@@ -1,80 +1,7 @@
-#!/usr/bin/env python3
+#!/usr/bin/env bash
-import argparse
-import json
-import os
+set -o errexit
+set -o nounset
-import boto3
-import tqdm
-
-from _common import ACCOUNT_NAMES, get_aws_session
-
-
-def get_session(*, queue_url):
- """
- Return a boto3 Session for publishing to SNS.
-
- If it recognises the account which contains the queue, it will pick
- the appropriate IAM role, otherwise it use the default boto3 Session.
- """
- # The arn format of an SQS queue URL is:
- #
- # https://sqs.eu-west-1.amazonaws.com/1234567890/queue-name
- #
- # Extract the account ID.
- account_id = queue_url.split("/")[3]
-
- try:
- role_arn = (
- f"arn:aws:iam::{account_id}:role/{ACCOUNT_NAMES[account_id]}-developer"
- )
- return get_aws_session(role_arn=role_arn)
- except KeyError:
- return boto3.Session()
-
-
-def parse_args():
- parser = argparse.ArgumentParser(
- prog=os.path.basename(__file__),
- description="Download all the messages from an SQS queue.",
- )
-
- parser.add_argument(
- "QUEUE_URL", help="The URL of the SQS queue to fetch messages from"
- )
-
- return parser.parse_args()
-
-
-def download_messages(*, queue_url):
- sess = get_session(queue_url=queue_url)
-
- sqs_client = sess.client("sqs")
-
- while True:
- resp = sqs_client.receive_message(
- QueueUrl=queue_url, AttributeNames=["All"], MaxNumberOfMessages=10
- )
-
- try:
- yield from resp["Messages"]
- except KeyError:
- return
-
- entries = [
- {"Id": msg["MessageId"], "ReceiptHandle": msg["ReceiptHandle"]}
- for msg in resp["Messages"]
- ]
-
- resp = sqs_client.delete_message_batch(QueueUrl=queue_url, Entries=entries)
-
- if len(resp["Successful"]) != len(entries):
- raise RuntimeError(
- f"Failed to delete messages: entries={entries!r} resp={resp!r}"
- )
-
-
-if __name__ == "__main__":
- args = parse_args()
- for message in tqdm.tqdm(download_messages(queue_url=args.QUEUE_URL)):
- print(json.dumps(message))
+_ensure_aws_credentials_are_fresh
+download_sqs_messages.py "$@"
aws/download_sqs_messages.py (0) → aws/download_sqs_messages.py (2121)
diff --git a/aws/download_sqs_messages.py b/aws/download_sqs_messages.py
new file mode 100755
index 0000000..201dcda
--- /dev/null
+++ b/aws/download_sqs_messages.py
@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+
+import json
+import os
+import sys
+
+import boto3
+import tqdm
+
+from _common import ACCOUNT_NAMES, get_aws_session
+
+
+def list_queue_urls():
+ sess = boto3.Session()
+ paginator = sess.client("sqs").get_paginator("list_queues")
+
+ for page in paginator.paginate():
+ yield from page["QueueUrls"]
+
+
+def get_session(*, queue_url):
+ """
+ Return a boto3 Session for publishing to SNS.
+
+ If it recognises the account which contains the queue, it will pick
+ the appropriate IAM role, otherwise it use the default boto3 Session.
+ """
+ # The arn format of an SQS queue URL is:
+ #
+ # https://sqs.eu-west-1.amazonaws.com/1234567890/queue-name
+ #
+ # Extract the account ID.
+ account_id = queue_url.split("/")[3]
+
+ try:
+ role_arn = (
+ f"arn:aws:iam::{account_id}:role/{ACCOUNT_NAMES[account_id]}-developer"
+ )
+ return get_aws_session(role_arn=role_arn)
+ except KeyError:
+ return boto3.Session()
+
+
+def download_messages(*, queue_url):
+ sess = get_session(queue_url=queue_url)
+
+ sqs_client = sess.client("sqs")
+
+ while True:
+ resp = sqs_client.receive_message(
+ QueueUrl=queue_url, AttributeNames=["All"], MaxNumberOfMessages=10
+ )
+
+ try:
+ yield from resp["Messages"]
+ except KeyError:
+ return
+
+ entries = [
+ {"Id": msg["MessageId"], "ReceiptHandle": msg["ReceiptHandle"]}
+ for msg in resp["Messages"]
+ ]
+
+ resp = sqs_client.delete_message_batch(QueueUrl=queue_url, Entries=entries)
+
+ if len(resp["Successful"]) != len(entries):
+ raise RuntimeError(
+ f"Failed to delete messages: entries={entries!r} resp={resp!r}"
+ )
+
+
+if __name__ == "__main__":
+ try:
+ queue_url = sys.argv[1]
+ except IndexError:
+ queue_url = None
+
+
+ if queue_url is not None:
+ for message in tqdm.tqdm(download_messages(queue_url=queue_url)):
+ print(json.dumps(message))
+ else:
+ for queue_url in list_queue_urls():
+ print(queue_url)
\ No newline at end of file