feat: support a mode to drain pulse queues rather than listen

src/fxci_etl/console.py

@@ -1,4 +1,3 @@
-import asyncio
 import sys
 from dataclasses import asdict
 from pathlib import Path
@@ -10,7 +9,7 @@
 
 from fxci_etl.config import Config
 from fxci_etl.metric.export import export_metrics
-from fxci_etl.pulse.listen import listen
+from fxci_etl.pulse.consume import drain, listen
 
 APP_NAME = "fxci-etl"
 
@@ -42,8 +41,26 @@ def handle(self):
             )
             return 1
 
-        loop = asyncio.get_event_loop()
-        loop.run_until_complete(listen(config, queue))
+        listen(config, queue)
+        return 0
+
+
+class PulseDrainCommand(ConfigCommand):
+    name = "pulse drain"
+    description = "Process events in the specified pulse queue and exit."
+    arguments = [argument("queue", description="Pulse queue to drain.")]
+
+    def handle(self):
+        config = self.parse_config(self.option("config"))
+        queue = self.argument("queue")
+
+        if queue not in config.pulse.queues:
+            self.line_error(
+                f"Queue '{queue}' is not configured! Choose from: {', '.join(config.pulse.queues)}"
+            )
+            return 1
+
+        drain(config, queue)
         return 0
 
 

src/fxci_etl/pulse/consume.py

@@ -0,0 +1,59 @@
+from kombu import Connection, Exchange, Queue
+
+from fxci_etl.config import Config
+from fxci_etl.pulse.handlers.base import handlers
+
+
+def get_connection(config: Config):
+    pulse = config.pulse
+    return Connection(
+        hostname=pulse.host,
+        port=pulse.port,
+        userid=pulse.user,
+        password=pulse.password,
+        ssl=True,
+    )
+
+
+def get_consumer(config: Config, connection: Connection, name: str):
+    pulse = config.pulse
+    qconf = pulse.queues[name]
+    exchange = Exchange(qconf.exchange, type="topic")
+    exchange(connection).declare(
+        passive=True
+    )  # raise an error if exchange doesn't exist
+
+    queue = Queue(
+        name=f"queue/{pulse.user}/{name}",
+        exchange=exchange,
+        routing_key=qconf.routing_key,
+        durable=qconf.durable,
+        exclusive=False,
+        auto_delete=qconf.auto_delete,
+    )
+
+    callbacks = [
+        cls(config)
+        for name, cls in handlers.items()
+        if config.etl.handlers is None or name in config.etl.handlers
+    ]
+    consumer = connection.Consumer(queue, auto_declare=False, callbacks=callbacks)
+    consumer.queues[0].queue_declare()
+    consumer.queues[0].queue_bind()
+    return consumer
+
+
+def listen(config: Config, name: str):
+    with get_connection(config) as connection:
+        with get_consumer(config, connection, name):
+            while True:
+                try:
+                    connection.drain_events(timeout=None)
+                except TimeoutError:
+                    pass
+
+
+def drain_events(config: Config, name: str):
+    with get_connection(config) as connection:
+        with get_consumer(config, connection, name):
+            connection.drain_events(timeout=0)