diff --git a/src/murfey/cli/repost_failed_calls.py b/src/murfey/cli/repost_failed_calls.py index e7f0c7f32..859547b59 100644 --- a/src/murfey/cli/repost_failed_calls.py +++ b/src/murfey/cli/repost_failed_calls.py @@ -1,6 +1,7 @@ import argparse import asyncio import json +import os from datetime import datetime from functools import partial from inspect import getfullargspec, iscoroutinefunction @@ -10,10 +11,6 @@ from sqlmodel import Session, create_engine from workflows.transport.pika_transport import PikaTransport -from murfey.server.murfey_db import url -from murfey.server.run import _set_up_transport -from murfey.util.config import security_from_file - def dlq_purge( dlq_dump_path: Path, queue: str, rabbitmq_credentials: Path @@ -191,10 +188,17 @@ def run(): args = parser.parse_args() # Read the security config file + os.environ["MURFEY_SECURITY_CONFIGURATION"] = args.config + + from murfey.util.config import security_from_file + security_config = security_from_file(args.config) # Configure the transport PikaTransport().load_configuration_file(security_config.rabbitmq_credentials) + + from murfey.server.run import _set_up_transport + _set_up_transport("PikaTransport") # Now import transport object which was set up in the above step @@ -212,6 +216,8 @@ def run(): ) # Set up database and retry api calls + from murfey.server.murfey_db import url + _url = url(security_config) engine = create_engine(_url) with Session(engine) as murfey_db: diff --git a/src/murfey/server/api/instrument.py b/src/murfey/server/api/instrument.py index 035563081..c04caeab9 100644 --- a/src/murfey/server/api/instrument.py +++ b/src/murfey/server/api/instrument.py @@ -3,7 +3,6 @@ import asyncio import datetime import logging -import os from pathlib import Path from typing import Annotated, Any, List, Optional from urllib.parse import quote @@ -16,20 +15,14 @@ try: from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_common.schemas import AcquisitionData, MicroscopeData - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False diff --git a/src/murfey/server/api/session_control.py b/src/murfey/server/api/session_control.py index bbc03fe91..70a240da8 100644 --- a/src/murfey/server/api/session_control.py +++ b/src/murfey/server/api/session_control.py @@ -1,4 +1,3 @@ -import os from datetime import datetime from logging import getLogger from pathlib import Path @@ -15,20 +14,14 @@ try: from smartem_agent.fs_parser import EpuParser from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_common.schemas import AtlasTileGridSquarePositionData - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False diff --git a/src/murfey/server/api/workflow.py b/src/murfey/server/api/workflow.py index a235a113c..85b2c06d2 100644 --- a/src/murfey/server/api/workflow.py +++ b/src/murfey/server/api/workflow.py @@ -1,5 +1,4 @@ import asyncio -import os from datetime import datetime from logging import getLogger from pathlib import Path @@ -24,7 +23,6 @@ try: from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_common.schemas import ( AcquisitionData as SmartEMAcquisitionData, GridData as SmartEMGridData, @@ -32,17 +30,12 @@ MicrographManifest as SmartEMMicrographManifest, ) - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False @@ -695,8 +688,6 @@ async def request_spa_preprocessing( "Failed to register micrograph with smartem", exc_info=True ) - db.close() - recipe_name = machine_config.recipes.get( "em-spa-preprocess", "em-spa-preprocess" ) @@ -741,7 +732,7 @@ async def request_spa_preprocessing( f"Pre-processing was requested for {sanitise(Path(proc_file.path).name)} " "but no Zocalo transport object was found" ) - return proc_file + db.close() else: for_stash = PreprocessStash( diff --git a/src/murfey/server/feedback.py b/src/murfey/server/feedback.py index bbd3fbf1a..d338e3f9f 100644 --- a/src/murfey/server/feedback.py +++ b/src/murfey/server/feedback.py @@ -52,6 +52,7 @@ logger = logging.getLogger("murfey.server.feedback") +murfey_entry_points = entry_points(group="murfey.workflows") # The first job number available to dynamic SPA feedback jobs. Jobs 1..6 are # the fixed preprocessing jobs (Import, MotionCorr, CtfFind, AutoPick, Extract, @@ -2270,11 +2271,11 @@ def feedback_callback(header: dict, message: dict, _db=murfey_db) -> None: if murfey.server._transport_object: murfey.server._transport_object.transport.ack(header) return None - elif message["register"] in entry_points(group="murfey.workflows").names: + elif message["register"] in murfey_entry_points.names: # Search for corresponding workflow - workflows: list[EntryPoint] = list( - entry_points(group="murfey.workflows", name=message["register"]) - ) # Returns either 1 item or empty list + workflows: list[EntryPoint] = [ + ep for ep in murfey_entry_points if ep.name == message["register"] + ] # Returns either 1 item or empty list if not workflows: logger.error(f"No workflow found for {sanitise(message['register'])}") if murfey.server._transport_object: diff --git a/src/murfey/util/config.py b/src/murfey/util/config.py index ed684ade5..564bbe505 100644 --- a/src/murfey/util/config.py +++ b/src/murfey/util/config.py @@ -12,6 +12,13 @@ from pydantic import BaseModel, ConfigDict, RootModel, ValidationInfo, field_validator from pydantic_settings import BaseSettings +try: + from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config + + SMARTEM_ACTIVE = True +except ImportError: + SMARTEM_ACTIVE = False + class MagnificationTable(RootModel[dict[int, float]]): pass @@ -353,3 +360,16 @@ def get_extended_machine_config( model = list(entry_points(group="murfey.config", name=extension_name))[0].load() data = getattr(machine_config, extension_name, {}) return model(**data) + + +def get_smartem_keycloak_client(): + keycloak_client = None + if SMARTEM_ACTIVE: + if keycloak_config := ( + os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") + or get_security_config().smartem_keycloak_config + ): + keycloak_client = KeycloakClient( + load_keycloak_config(Path(keycloak_config)) + ) + return keycloak_client diff --git a/src/murfey/util/processing_params.py b/src/murfey/util/processing_params.py index 6cf1fbc46..918106775 100644 --- a/src/murfey/util/processing_params.py +++ b/src/murfey/util/processing_params.py @@ -79,8 +79,9 @@ def motion_corrected_mrc( else: parts = [secure_filename(p) for p in input_movie.parts] visit_idx = parts.index(visit_name) - core = Path(*parts[: visit_idx + 1]) - ppath = Path(*parts) + # Include leading slash which is removed in the secure parts above + core = Path("/") / Path(*parts[: visit_idx + 1]) + ppath = Path("/") / Path(*parts) if machine_config.process_multiple_datasets: sub_dataset = ppath.relative_to(core).parts[0] else: diff --git a/src/murfey/workflows/spa/ctf_estimation.py b/src/murfey/workflows/spa/ctf_estimation.py index 799cc0612..7e78374ce 100644 --- a/src/murfey/workflows/spa/ctf_estimation.py +++ b/src/murfey/workflows/spa/ctf_estimation.py @@ -1,6 +1,4 @@ -import os from logging import getLogger -from pathlib import Path from sqlmodel import Session, select @@ -14,7 +12,6 @@ try: from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_backend.model.http_request import ( CtfEstimationRegisteredRequest, MicrographUpdateRequest, @@ -25,17 +22,12 @@ ) from smartem_common.entity_status import MicrographStatus - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False diff --git a/src/murfey/workflows/spa/flush_spa_preprocess.py b/src/murfey/workflows/spa/flush_spa_preprocess.py index 81363ca50..f63700d31 100644 --- a/src/murfey/workflows/spa/flush_spa_preprocess.py +++ b/src/murfey/workflows/spa/flush_spa_preprocess.py @@ -1,5 +1,4 @@ import logging -import os from pathlib import Path from typing import Optional @@ -10,24 +9,18 @@ try: from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_common.schemas import ( FoilHoleData as SmartEMFoilHoleData, GridSquareData as SmartEMGridSquareData, GridSquareMetadata as SmartEMGridSquareMetadata, ) - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False diff --git a/src/murfey/workflows/spa/motion_correction.py b/src/murfey/workflows/spa/motion_correction.py index 218da437b..d1b0c5043 100644 --- a/src/murfey/workflows/spa/motion_correction.py +++ b/src/murfey/workflows/spa/motion_correction.py @@ -1,6 +1,4 @@ -import os from logging import getLogger -from pathlib import Path from sqlmodel import Session, select @@ -14,7 +12,6 @@ try: from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_backend.model.http_request import ( MicrographUpdateRequest, MotionCorrectionRegisteredRequest, @@ -25,17 +22,12 @@ ) from smartem_common.entity_status import MicrographStatus - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False diff --git a/src/murfey/workflows/spa/picking.py b/src/murfey/workflows/spa/picking.py index a2000a55b..7de4bcbce 100644 --- a/src/murfey/workflows/spa/picking.py +++ b/src/murfey/workflows/spa/picking.py @@ -1,6 +1,4 @@ -import os from logging import getLogger -from pathlib import Path from typing import List import numpy as np @@ -33,22 +31,16 @@ try: from smartem_backend.api_client import SmartEMAPIClient - from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_backend.model.http_request import MicrographUpdateRequest from smartem_backend.model.http_response import MicrographResponse from smartem_common.entity_status import MicrographStatus - from murfey.util.config import get_security_config + from murfey.util.config import get_smartem_keycloak_client - keycloak_client = KeycloakClient( - load_keycloak_config( - Path( - os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") - or get_security_config().smartem_keycloak_config - ) - ) - ) - SMARTEM_ACTIVE = True + if keycloak_client := get_smartem_keycloak_client(): + SMARTEM_ACTIVE = True + else: + SMARTEM_ACTIVE = False except ImportError: keycloak_client = None SMARTEM_ACTIVE = False diff --git a/tests/cli/test_repost_failed_calls.py b/tests/cli/test_repost_failed_calls.py index 1bb841d52..844665826 100644 --- a/tests/cli/test_repost_failed_calls.py +++ b/tests/cli/test_repost_failed_calls.py @@ -192,7 +192,7 @@ def test_handle_failed_posts(tmp_path): @mock.patch("murfey.cli.repost_failed_calls.dlq_purge") @mock.patch("murfey.cli.repost_failed_calls.handle_failed_posts") @mock.patch("murfey.cli.repost_failed_calls.handle_dlq_messages") -@mock.patch("murfey.cli.repost_failed_calls.url") +@mock.patch("murfey.server.murfey_db.url") @mock.patch("murfey.cli.repost_failed_calls.create_engine") @mock.patch("murfey.cli.repost_failed_calls.Session") def test_run_repost_failed_calls(