Simplification is possible as don't need to store messages

Author: stephen-riggs

src/murfey/cli/dlq_resubmit.py

@@ -1,7 +1,6 @@
 import argparse
 import json
 import os
-import time
 from datetime import datetime
 from functools import partial
 from pathlib import Path
@@ -11,10 +10,10 @@
 from jose import jwt
 from workflows.transport.pika_transport import PikaTransport
 
-dlq_dump_path = Path("./DLQ")
 
-
-def dlq_purge(queue: str, rabbitmq_credentials: Path) -> list[Path]:
+def dlq_purge(
+    dlq_dump_path: Path, queue: str, rabbitmq_credentials: Path
+) -> list[Path]:
     transport = PikaTransport()
     transport.load_configuration_file(rabbitmq_credentials)
     transport.connect()
@@ -26,26 +25,8 @@ def dlq_purge(queue: str, rabbitmq_credentials: Path) -> list[Path]:
     def receive_dlq_message(header: dict, message: dict) -> None:
         idlequeue.put_nowait("start")
         header["x-death"][0]["time"] = datetime.timestamp(header["x-death"][0]["time"])
-
-        timestamp = time.localtime(int(header["x-death"][0]["time"]))
-        filepath = dlq_dump_path / time.strftime("%Y-%m-%d", timestamp)
-        filepath.mkdir(parents=True, exist_ok=True)
-        filename = filepath / (
-            f"{queue}-"
-            + time.strftime("%Y%m%d-%H%M%S", timestamp)
-            + "-"
-            + str(header["message-id"])
-        )
-
-        dlqmsg = {
-            "exported": {
-                "date": time.strftime("%Y-%m-%d"),
-                "time": time.strftime("%H:%M:%S"),
-            },
-            "header": header,
-            "message": message,
-        }
-
+        filename = dlq_dump_path / (f"{queue}-" + str(header["message-id"]))
+        dlqmsg = {"header": header, "message": message}
         with filename.open("w") as fh:
             json.dump(dlqmsg, fh, indent=2, sort_keys=True)
         print(f"Message {header['message-id']} exported to {filename}")
@@ -60,11 +41,10 @@ def receive_dlq_message(header: dict, message: dict) -> None:
         acknowledgement=True,
     )
     try:
-        idlequeue.get(True, 3)
         while True:
             idlequeue.get(True, 0.1)
     except Empty:
-        print("Done.")
+        print("Done dlq purge")
     transport.disconnect()
     return exported_messages
 
@@ -75,20 +55,10 @@ def handle_dlq_messages(messages_path: list[Path], rabbitmq_credentials: Path):
     transport.connect()
 
     for f, dlqfile in enumerate(messages_path):
-        if not Path(dlqfile).is_file():
-            print(f"Ignoring missing file {dlqfile}")
+        if not dlqfile.is_file():
             continue
         with open(dlqfile) as fh:
             dlqmsg = json.load(fh)
-        print(f"Parsing message from {dlqfile}")
-        if (
-            not isinstance(dlqmsg, dict)
-            or not dlqmsg.get("header")
-            or not dlqmsg.get("message")
-        ):
-            print(f"File {dlqfile} is not a valid DLQ message.")
-            continue
-
         header = dlqmsg["header"]
         header["dlq-reinjected"] = "True"
 
@@ -109,7 +79,7 @@ def handle_dlq_messages(messages_path: list[Path], rabbitmq_credentials: Path):
             headers=clean_header,
         )
         dlqfile.unlink()
-        print(f"Done {dlqfile}\n")
+        print(f"Reinjected {dlqfile}\n")
 
     transport.disconnect()
 
@@ -158,6 +128,9 @@ def run():
         help="Token username",
         required=True,
     )
+    parser.add_argument(
+        "-d", "--dir", default="DLQ", help="Directory to export messages to"
+    )
     args = parser.parse_args()
 
     # Set the environment variable then read it by importing the security config
@@ -174,18 +147,17 @@ def run():
     )
 
     # Purge the queue and repost/reinject any messages found
+    dlq_dump_path = Path(args.dir)
+    dlq_dump_path.mkdir(parents=True, exist_ok=True)
     exported_messages = dlq_purge(
-        security_config.feedback_queue, security_config.rabbitmq_credentials
+        dlq_dump_path,
+        security_config.feedback_queue,
+        security_config.rabbitmq_credentials,
     )
     handle_failed_posts(exported_messages, token)
     handle_dlq_messages(exported_messages, security_config.rabbitmq_credentials)
 
     # Clean up any created directories
-    for date_directory in dlq_dump_path.glob("*"):
-        try:
-            date_directory.rmdir()
-        except OSError:
-            print(f"Cannot remove {date_directory} as it is not empty")
     try:
         dlq_dump_path.rmdir()
     except OSError: