ID

d52e547

date

2022-10-04 09:42:32+00:00

author

Alex Chan <alex@alexwlchan.net>

parent

b2261c5

message

Add an SQS stats script

changed files

1 file, 96 additions

• sqs_stats

Changed files

sqs_stats (0) → sqs_stats (2486)

diff --git a/sqs_stats b/sqs_stats
new file mode 100755
index 0000000..0152208
--- /dev/null
+++ b/sqs_stats
@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+
+import collections
+import sys
+
+import boto3
+import humanize
+import termcolor
+
+
+# See https://github.com/boto/botocore/issues/2705
+import warnings
+warnings.filterwarnings('ignore', category=FutureWarning, module='botocore.client')
+
+
+
+def list_queue_urls_in_account(sess, *, prefixes):
+    """
+    Generates a list of all the queue URLs in an account.
+    """
+    sqs_client = sess.client("sqs")
+
+    for prefix in prefixes:
+        for page in sqs_client.get_paginator("list_queues").paginate(
+            QueueNamePrefix=prefix
+        ):
+            yield from page["QueueUrls"]
+
+
+def get_queue_stats(sess, *, queue_urls):
+    """
+    Get the size of the queues associated with this pipeline.
+    """
+    sqs_client = sess.client("sqs")
+
+    attribute_names = [
+        "ApproximateNumberOfMessages",
+        "ApproximateNumberOfMessagesNotVisible",
+        "ApproximateNumberOfMessagesDelayed",
+    ]
+
+    queue_responses = {
+        q_url: sqs_client.get_queue_attributes(
+            QueueUrl=q_url, AttributeNames=attribute_names
+        )
+        for q_url in queue_urls
+    }
+
+    return {
+        q_url: sum(int(resp["Attributes"][attr]) for attr in attribute_names)
+        for q_url, resp in queue_responses.items()
+    }
+
+
+def print_number(value, color):
+    if value is None:
+        print("-".rjust(9, " "), end="")
+    else:
+        print(termcolor.colored(humanize.intcomma(value).rjust(9, " "), color), end="")
+
+
+def pprint_queue_stats(queue_stats):
+    interesting_queues = {
+        q_url: q_size for q_url, q_size in queue_stats.items() if q_size > 0
+    }
+
+    if not interesting_queues:
+        print("All queues are empty")
+        return
+
+    paired_queues = collections.defaultdict(lambda: {"q": None, "dlq": None})
+
+    for q_url, q_size in interesting_queues.items():
+        q_name = q_url.split("/")[-1]
+        if q_name.endswith("_dlq"):
+            paired_queues[q_name.replace("_dlq", "")]["dlq"] = q_size
+        else:
+            paired_queues[q_name]["q"] = q_size
+
+    for q_name, q_stats in sorted(paired_queues.items()):
+        print_number(q_stats["q"], "green")
+        print_number(q_stats["dlq"], "red")
+        print("\t", end="")
+        print(q_name)
+
+
+if __name__ == "__main__":
+    sess = boto3.Session()
+
+    prefixes = sys.argv[1:] or ("",)
+
+    queue_urls = list_queue_urls_in_account(sess, prefixes=prefixes)
+
+    queue_stats = get_queue_stats(sess, queue_urls=queue_urls)
+
+    pprint_queue_stats(queue_stats)