Refactored amqp_helpers.py to use a global connection object (I was

originally going to use a singleton pattern but ran into issues with
that)

apps/cli/utilities/wf_monitor/src/wf_monitor/monitor.py

@@ -13,7 +13,7 @@ from pathlib import Path
 import pendulum
 
 from ._version import ___version___ as VERSION
-from channels.amqp_helpers import make_amqp_connection, configure_amqp_channel, emit_amqp_message
+from channels.amqp_helpers import workflow_events
 from workspaces.schema import WorkflowEvent, WorkflowEventType
 
 WORKFLOW_STATUS_EXCH = 'workspaces.workflow-service.workflow-status'
@@ -239,23 +239,26 @@ def main():
     # Parse command-line args
     args = make_arg_parser().parse_args()
     monitor = WorkflowMonitor(args.log_path)
-
-    # Probably want to refactor this so that it doesn't create its own connection just for this,
-    # since creating connections is expensive
-    decorated_connect = log_decorator_factory('Making connection...')(make_amqp_connection)
-    decorated_config_channel = log_decorator_factory('Configuring channel...')(configure_amqp_channel)
-    decorated_emit = log_decorator_factory('Emitting message...')(emit_amqp_message)
-
-    with decorated_connect(args.profile) as connection:
-        with decorated_config_channel(
-            connection=connection,
-            exchange=WORKFLOW_STATUS_EXCH,
-            exchange_type='topic'
-        ) as channel:
-            for event in monitor.events:
-                decorated_emit(
-                    channel=channel,
-                    exchange=WORKFLOW_STATUS_EXCH,
-                    routing_key=f'{event.job_name}.{event.job_id}.{event.type.name.lower()}',
-                    msg=event.json()
-                )
+    with workflow_events as wfe:
+        for event in monitor.events:
+            wfe.send(event)
+
+    # # Probably want to refactor this so that it doesn't create its own connection just for this,
+    # # since creating connections is expensive
+    # decorated_connect = log_decorator_factory('Making connection...')(make_amqp_connection)
+    # decorated_config_channel = log_decorator_factory('Configuring channel...')(configure_amqp_channel)
+    # decorated_emit = log_decorator_factory('Emitting message...')(emit_amqp_message)
+    #
+    # with decorated_connect(args.profile) as connection:
+    #     with decorated_config_channel(
+    #         connection=connection,
+    #         exchange=WORKFLOW_STATUS_EXCH,
+    #         exchange_type='topic'
+    #     ) as channel:
+    #         for event in monitor.events:
+    #             decorated_emit(
+    #                 channel=channel,
+    #                 exchange=WORKFLOW_STATUS_EXCH,
+    #                 routing_key=f'{event.job_name}.{event.job_id}.{event.type.name.lower()}',
+    #                 msg=event.json()
+    #             )

shared/channels/src/channels/amqp_helpers.py

@@ -9,6 +9,7 @@ from workspaces.json import WorkflowEventSchema
 from workspaces.schema import WorkflowEvent
 
 T = TypeVar('T', contravariant=True)
+CONN = None
 
 
 class ChannelDefinition(Protocol[T]):
@@ -53,12 +54,10 @@ ChannelDef = TypeVar('ChannelDef', bound=ChannelDefinition, covariant=True)
 
 
 class Channel(Generic[ChannelDef]):
-    CONNECTION : pika.BlockingConnection = None
-    CHANNEL : BlockingChannel = None
-    CONFIG : CapoConfig = None
-
     def __init__(self, definition: ChannelDefinition[T]):
-        self.definition = definition
+        self.definition: ChannelDefinition[T] = definition
+        self.chan: BlockingChannel = None
+        self.config: CapoConfig = None
 
     def connect(self, **kwargs: Union[int, str]):
         """
@@ -74,32 +73,34 @@ class Channel(Generic[ChannelDef]):
         :param: password: Password to use when connecting
         :param: exchange: Exchange to use when connection
         """
-        if not self.CONNECTION:
-            self.CONFIG = CapoConfig(
+        global CONN
+
+        if not CONN:
+            self.config = CapoConfig(
                 profile=kwargs.get('profile', None)
             ).settings('edu.nrao.archive.configuration.AmqpServer')
 
             connection_parameters = pika.ConnectionParameters(
-                host=kwargs.get('hostname', self.CONFIG.hostname),
-                port=int(kwargs.get('port', self.CONFIG.port)),
+                host=kwargs.get('hostname', self.config.hostname),
+                port=int(kwargs.get('port', self.config.port)),
                 connection_attempts=int(kwargs.get('connection_attempts', 5)),
                 socket_timeout=int(kwargs.get('socket_timeout', 5000)),
                 retry_delay=int(kwargs.get('retry_delay', 500)),
                 credentials=pika.PlainCredentials(
-                    username=kwargs.get('username', self.CONFIG.username),
-                    password=kwargs.get('password', self.CONFIG.password)
+                    username=kwargs.get('username', self.config.username),
+                    password=kwargs.get('password', self.config.password)
                 )
             )
-            self.CONNECTION = pika.BlockingConnection(connection_parameters)
-            self.CHANNEL = self.CONNECTION.channel()
-            self.definition.declarations(self.CHANNEL)
+            CONN = pika.BlockingConnection(connection_parameters)
+        self.chan = CONN.channel()
+        self.definition.declarations(self.chan)
 
     def close(self):
         """
-        Close connection to RabbitMQ server, if it's open
+        Close channel of connection to RabbitMQ server, if it's open
         """
-        if self.CONNECTION:
-            self.CONNECTION.close()
+        if self.chan:
+            self.chan.close()
 
     def send(self, event: WorkflowEventSchema):
         """
@@ -108,7 +109,7 @@ class Channel(Generic[ChannelDef]):
         """
         rendered = self.definition.schema.dump(event)
         routing_key = self.definition.routing_key_for(event)
-        self.CHANNEL.basic_publish(routing_key=routing_key, body=rendered)
+        self.chan.basic_publish(routing_key=routing_key, body=rendered)
 
     def listen(self, callback: Optional[Callable], pattern: str = '#', auto_ack: bool = False):
         """
@@ -123,10 +124,10 @@ class Channel(Generic[ChannelDef]):
             event = self.definition.schema.load(message)
             callback(event)
 
-        queue = self.CHANNEL.queue_declare(queue='', exclusive=True).method.queue
-        self.CHANNEL.queue_bind(queue=queue, exchange=self.definition.exchange(), routing_key=pattern)
-        self.CHANNEL.basic_consume(queue=queue, on_message_callback=callback, auto_ack=auto_ack)
-        self.CHANNEL.start_consuming()
+        queue = self.chan.queue_declare(queue='', exclusive=True).method.queue
+        self.chan.queue_bind(queue=queue, exchange=self.definition.exchange(), routing_key=pattern)
+        self.chan.basic_consume(queue=queue, on_message_callback=callback, auto_ack=auto_ack)
+        self.chan.start_consuming()
 
     def __enter__(self, **kwargs: Union[int, str]):
         """
@@ -141,10 +142,11 @@ class Channel(Generic[ChannelDef]):
         :param: exchange: Exchange to use when connection
         """
         self.connect(**kwargs)
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
 
 # Predefined event channels for ease of use
-workflow_events : Channel[WorkflowEvent] = Channel(WorkflowEventChannel())
+workflow_events: Channel[WorkflowEvent] = Channel(WorkflowEventChannel())