Load RabbitMQ credentials directly using PikaTransport object (#434)

* Use PikaTransport object directly to load RabbitMQ credentials
* Ignore unknown arguments when starting up Murfey server from command line
* Fixed demo mode need RabbitMQ credentials

---------

Co-authored-by: Dan Hatton <[email protected]>

Author: tieneupin

src/murfey/server/__init__.py

@@ -11,13 +11,12 @@
 from importlib.resources import files
 from pathlib import Path
 from threading import Thread
-from typing import Any, Dict, List, NamedTuple, Tuple
+from typing import Any, Dict, List, Literal, NamedTuple, Tuple
 
 import graypy
 import mrcfile
 import numpy as np
 import uvicorn
-import workflows
 from backports.entry_points_selectable import entry_points
 from fastapi import Request
 from fastapi.templating import Jinja2Templates
@@ -42,6 +41,7 @@
 from sqlalchemy.orm.exc import ObjectDeletedError
 from sqlmodel import Session, create_engine, select
 from werkzeug.utils import secure_filename
+from workflows.transport.pika_transport import PikaTransport
 
 import murfey
 import murfey.server.ispyb
@@ -225,6 +225,7 @@ def respond_with_template(
 
 
 def run():
+    # Set up argument parser
     parser = argparse.ArgumentParser(description="Start the Murfey server")
     parser.add_argument(
         "--host",
@@ -273,28 +274,29 @@ def run():
         help="Increase logging output verbosity",
         default=0,
     )
+    # Parse and separate known and unknown args
+    args, unknown = parser.parse_known_args()
 
+    # Load the security configuration
     security_config = get_security_config()
-    # setup logging
+
+    # Set up GrayLog handler if provided in the configuration
     if security_config.graylog_host:
         handler = graypy.GELFUDPHandler(
             security_config.graylog_host, security_config.graylog_port, level_names=True
         )
         root_logger = logging.getLogger()
         root_logger.addHandler(handler)
-
     # Install a log filter to all existing handlers.
     LogFilter.install()
 
-    workflows.transport.load_configuration_file(security_config.rabbitmq_credentials)
-
-    args = parser.parse_args()
-
-    # Set up Zocalo connection
     if args.demo:
+        # Run in demo mode with no connections set up
         os.environ["MURFEY_DEMO"] = "1"
     else:
-        _set_up_transport(args.transport)
+        # Load RabbitMQ configuration and set up the connection
+        PikaTransport().load_configuration_file(security_config.rabbitmq_credentials)
+        _set_up_transport("PikaTransport")
 
     # Set up logging now that the desired verbosity is known
     _set_up_logging(quiet=args.quiet, verbosity=args.verbose)
@@ -393,7 +395,7 @@ def _set_up_logging(quiet: bool, verbosity: int):
         logging.getLogger(logger_name).setLevel(log_level)
 
 
-def _set_up_transport(transport_type):
+def _set_up_transport(transport_type: Literal["PikaTransport"]):
     global _transport_object
     _transport_object = TransportManager(transport_type)
 

src/murfey/server/ispyb.py

@@ -3,7 +3,7 @@
 import datetime
 import logging
 import os
-from typing import Callable, List, Optional
+from typing import Callable, List, Literal, Optional
 
 import ispyb
 
@@ -55,7 +55,7 @@ def _send_using_new_connection(transport_type: str, queue: str, message: dict) -
 
 
 class TransportManager:
-    def __init__(self, transport_type):
+    def __init__(self, transport_type: Literal["PikaTransport"]):
         self._transport_type = transport_type
         self.transport = workflows.transport.lookup(transport_type)()
         self.transport.connect()