diff --git a/src/murfey/server/__init__.py b/src/murfey/server/__init__.py index c5a2825b..fb8e1a1c 100644 --- a/src/murfey/server/__init__.py +++ b/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) diff --git a/src/murfey/server/ispyb.py b/src/murfey/server/ispyb.py index ce01b2cd..119cbb7e 100644 --- a/src/murfey/server/ispyb.py +++ b/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()