New zocalo.pickup command (#223)

* New zocalo.pickup command

This reads dropfiles from the zocalo.go.fallback_location directory
which is where zocalo.go writes messages if the message broker is
unavailable. When the message broker becomes available again,
zocalo.pickup can be used to re-submit messages to the processing_recipe
queue.

* Sort dropfiles by mtime

Author: rjgildea

HISTORY.rst

@@ -5,6 +5,7 @@ History
 Unreleased
 ----------
 * Add Dockerfile and build-and-push-docker-image GitHub workflow
+* Add ``zocalo.pickup`` command for re-submitting messages stored in the ``zocalo.go.fallback_location`` while the message broker is unavailable
 
 0.26.0 (2022-11-04)
 -------------------

setup.cfg

@@ -47,6 +47,7 @@ console_scripts =
     zocalo.dlq_purge = zocalo.cli.dlq_purge:run
     zocalo.dlq_reinject = zocalo.cli.dlq_reinject:run
     zocalo.go = zocalo.cli.go:run
+    zocalo.pickup = zocalo.cli.pickup:run
     zocalo.queue_drain = zocalo.cli.queue_drain:run
     zocalo.service = zocalo.service:start_service
     zocalo.shutdown = zocalo.cli.shutdown:run

src/zocalo/cli/pickup.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import argparse
+import json
+import pathlib
+import sys
+import time
+
+import workflows.transport
+
+import zocalo.configuration
+
+
+def run():
+    zc = zocalo.configuration.from_file()
+    zc.activate()
+    dropdir = pathlib.Path(zc.storage["zocalo.go.fallback_location"])
+
+    parser = argparse.ArgumentParser(
+        usage="zocalo.pickup [options]", description="Processes zocalo.go backlog"
+    )
+
+    parser.add_argument("-?", action="help", help=argparse.SUPPRESS)
+    parser.add_argument(
+        "-d",
+        "--delay",
+        dest="delay",
+        action="store",
+        type=int,
+        default=2,
+        help="Number of seconds to wait between message dispatches",
+    )
+    parser.add_argument(
+        "-w",
+        "--wait",
+        dest="wait",
+        action="store",
+        type=int,
+        default=60,
+        help="Number of seconds to wait initially",
+    )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        dest="verbose",
+        action="store_true",
+        default=False,
+        help="Show raw message before sending",
+    )
+    zc.add_command_line_options(parser)
+    workflows.transport.add_command_line_options(parser, transport_argument=True)
+    args = parser.parse_args()
+
+    try:
+        files = list(dropdir.iterdir())
+    except OSError:
+        sys.exit("This program is only available to privileged users")
+
+    print(f"Found {len(files)} files")
+    if not files:
+        sys.exit()
+
+    if args.wait:
+        print(f"Waiting {args.wait} seconds")
+        time.sleep(args.wait)
+
+    print(f"Connecting to {args.transport}...")
+    transport = workflows.transport.lookup(args.transport)()
+    transport.connect()
+
+    file_info = {f: {} for f in files}
+
+    for f, finfo in file_info.items():
+        with f.open() as fh:
+            data = json.load(fh)
+            finfo["message"] = data["message"]
+            finfo["headers"] = data["headers"]
+        finfo["originating-host"] = finfo["headers"].get("zocalo.go.host")
+        finfo["recipes"] = ",".join(finfo["message"].get("recipes", []))
+        finfo["mtime"] = f.stat().st_mtime
+
+    count = 0
+    file_count = len(file_info)
+    for f, finfo in dict(
+        sorted(file_info.items(), key=lambda item: item[1]["mtime"])
+    ).items():
+        print(
+            f"Sending {f} from host {file_info[f]['originating-host']}"
+            f" with recipes {file_info[f]['recipes']}"
+        )
+        assert f.exists()
+        transport.send(
+            "processing_recipe",
+            file_info[f]["message"],
+            headers=file_info[f]["headers"],
+        )
+        f.unlink()
+        count = count + 1
+        print(f"Done ({count} of {file_count})")
+        try:
+            time.sleep(args.delay)
+        except KeyboardInterrupt:
+            print("CTRL+C - stopping")
+            time.sleep(0.5)
+            sys.exit(1)
+
+    transport.disconnect()

tests/cli/test_pickup.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import json
+import sys
+import time
+import uuid
+from unittest import mock
+
+import pytest
+import workflows.transport
+from workflows.transport.common_transport import CommonTransport
+
+import zocalo.configuration
+from zocalo.cli import pickup
+
+
+@pytest.fixture
+def mock_zocalo_configuration(tmp_path):
+    mock_zc = mock.MagicMock(zocalo.configuration.Configuration)
+    mock_zc.storage = {
+        "zocalo.go.fallback_location": str(tmp_path),
+    }
+    return mock_zc
+
+
+def test_pickup_empty_filelist_raises_system_exit(mocker, mock_zocalo_configuration):
+    mocker.patch.object(
+        zocalo.configuration, "from_file", return_value=mock_zocalo_configuration
+    )
+    with mock.patch.object(sys, "argv", ["prog"]), pytest.raises(SystemExit) as e:
+        pickup.run()
+        assert e.code == 0
+
+
+def test_pickup_sends_to_processing_recipe(mocker, mock_zocalo_configuration, tmp_path):
+    mocked_transport = mocker.MagicMock(CommonTransport)
+    mocker.patch.object(workflows.transport, "lookup", return_value=mocked_transport)
+    mocker.patch.object(
+        zocalo.configuration, "from_file", return_value=mock_zocalo_configuration
+    )
+    for i in range(10):
+        msg = {
+            "headers": {"zocalo.go.user": "foobar", "zocalo.go.host": "example.com"},
+            "message": {
+                "recipes": [f"thing{i}"],
+                "parameters": {"foo": i},
+            },
+        }
+        (tmp_path / str(uuid.uuid4())).write_text(json.dumps(msg))
+        time.sleep(0.1)
+    with mock.patch.object(sys, "argv", ["prog", "--wait", "0", "--delay", "0"]):
+        pickup.run()
+    mocked_transport().send.assert_has_calls(
+        [
+            mock.call(
+                "processing_recipe",
+                {
+                    "recipes": [f"thing{i}"],
+                    "parameters": {"foo": i},
+                },
+                headers=mock.ANY,
+            )
+            for i in range(10)
+        ]
+    )