diff --git a/src/murfey/server/api/instrument.py b/src/murfey/server/api/instrument.py index 9e7be9f5f..9606a0f6b 100644 --- a/src/murfey/server/api/instrument.py +++ b/src/murfey/server/api/instrument.py @@ -3,6 +3,7 @@ import asyncio import datetime import logging +import os from pathlib import Path from typing import Annotated, Any, List, Optional from urllib.parse import quote @@ -14,11 +15,16 @@ from werkzeug.utils import secure_filename try: - from smartem_backend.api_client import EntityConverter + from smartem_backend.api_client import SmartEMAPIClient + from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_common.schemas import AcquisitionData, MicroscopeData + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False import murfey.server.prometheus as prom @@ -156,32 +162,47 @@ async def setup_multigrid_watcher( session = db.exec(select(Session).where(Session.id == session_id)).one() visit = session.visit async with aiohttp.ClientSession() as clientsession: - acquisition_uuid = None - if SMARTEM_ACTIVE and machine_config.smartem_api_url: + acquisition_uuid = session.smartem_acquisition_uuid + if ( + SMARTEM_ACTIVE + and machine_config.smartem_api_url + and acquisition_uuid is None + ): log.info("registering an acquisition with smartem") try: microscope_data = MicroscopeData(instrument_id=instrument_name) - acquisition_data = EntityConverter.acquisition_to_request( - AcquisitionData( - name=visit, - id=visit, - instrument=microscope_data, - storage_path=str(secure_path(watcher_spec.source / visit)), - start_time=datetime.datetime.now(), - ) + smartem_client = SmartEMAPIClient( + base_url=machine_config.smartem_api_url, + logger=log, + keycloak_client=keycloak_client, + ) + acquisition_data = AcquisitionData( + name=visit, + id=visit, + instrument=microscope_data, + storage_path=str(secure_path(watcher_spec.source / visit)), + start_time=datetime.datetime.now(), + ) + acquisition_response_data = smartem_client.create_acquisition( + acquisition_data ) - async with clientsession.post( - f"{machine_config.smartem_api_url}/acquisitions", - json=acquisition_data.model_dump(mode="json"), - ) as response: - acquisition_response_data = await response.json() - acquisition_uuid = acquisition_response_data["uuid"] + acquisition_uuid = acquisition_response_data.uuid except Exception: log.warning( "failed to register acquisition with smartem", exc_info=True ) else: log.info("smartem not configured") + if acquisition_uuid is not None: + async with clientsession.post( + f"{machine_config.instrument_server_url}{url_path_for('api.router', 'update_session', session_id=session_id)}", + parameters={"smartem_acquisition_uuid": acquisition_uuid}, + headers={ + "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}" + }, + ) as resp: + await resp.json() + async with clientsession.post( f"{machine_config.instrument_server_url}{url_path_for('api.router', 'setup_multigrid_watcher', session_id=session_id)}", json={ diff --git a/src/murfey/server/api/session_control.py b/src/murfey/server/api/session_control.py index 703f62b90..c789b1901 100644 --- a/src/murfey/server/api/session_control.py +++ b/src/murfey/server/api/session_control.py @@ -1,3 +1,4 @@ +import os from datetime import datetime from logging import getLogger from pathlib import Path @@ -13,10 +14,15 @@ try: from smartem_backend.api_client import SmartEMAPIClient + from smartem_backend.keycloak_client import KeycloakClient, load_keycloak_config from smartem_common.schemas import AtlasData + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False import murfey.server.prometheus as prom @@ -379,7 +385,9 @@ def register_atlas( ] if machine_config.smartem_api_url: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) grid_uuid = None if atlas_registration_data.tag: @@ -398,6 +406,14 @@ def register_atlas( grid_uuid = grid.uuid break if grid_uuid is not None: + existing_atlas = smartem_client.get_grid_atlas(grid_uuid) + if ( + existing_atlas.name == atlas_registration_data.name + and existing_atlas.storage_folder + == atlas_registration_data.storage_folder + ): + # there is a question here of whether the grid should be registered if specified + return atlas_data = AtlasData( id=atlas_registration_data.name, acquisition_date=datetime.now(), @@ -449,7 +465,9 @@ def register_square( ] if machine_config.smartem_api_url: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) smartem_client.gridsquare_registered( smartem_uuid, count=square_registration_data.count diff --git a/src/murfey/server/api/session_info.py b/src/murfey/server/api/session_info.py index a28746e12..23f0e1bb0 100644 --- a/src/murfey/server/api/session_info.py +++ b/src/murfey/server/api/session_info.py @@ -196,10 +196,14 @@ def create_session( @router.post("/sessions/{session_id}") def update_session( - session_id: MurfeySessionID, process: bool = True, db=murfey_db + session_id: MurfeySessionID, + process: bool = True, + smartem_acquisition_uuid: str | None = None, + db=murfey_db, ) -> None: session = db.exec(select(Session).where(Session.id == session_id)).one() session.process = process + session.smartem_acquisition_uuid = smartem_acquisition_uuid db.add(session) db.commit() return None diff --git a/src/murfey/server/api/workflow.py b/src/murfey/server/api/workflow.py index a91dcd3ee..195897d3e 100644 --- a/src/murfey/server/api/workflow.py +++ b/src/murfey/server/api/workflow.py @@ -1,4 +1,5 @@ import asyncio +import os from datetime import datetime from logging import getLogger from pathlib import Path @@ -21,6 +22,7 @@ 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, @@ -28,8 +30,12 @@ MicrographManifest as SmartEMMicrographManifest, ) + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False import murfey.server.prometheus as prom @@ -134,7 +140,9 @@ def register_dc_group( if machine_config.smartem_api_url: try: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) grid_data = SmartEMGridData( data_dir=Path(dcg_params.tag), @@ -621,7 +629,9 @@ async def request_spa_preprocessing( fh, gs = fh_with_gs if fh.smartem_uuid: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) movie_path = Path(proc_file.path) micrograph_manifest = SmartEMMicrographManifest( diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py index 189f98e06..b8025a203 100644 --- a/src/murfey/util/db.py +++ b/src/murfey/util/db.py @@ -53,6 +53,7 @@ class Session(SQLModel, table=True): # type: ignore instrument_name: str = Field(default="") process: bool = Field(default=True) visit_end_time: Optional[datetime] = Field(default=None) + smartem_acquisition_uuid: str | None = Field(default=None) # Image sites associated with this session imaging_sites: List["ImagingSite"] = Relationship( diff --git a/src/murfey/workflows/spa/ctf_estimation.py b/src/murfey/workflows/spa/ctf_estimation.py index fbe6eda85..5f88abd77 100644 --- a/src/murfey/workflows/spa/ctf_estimation.py +++ b/src/murfey/workflows/spa/ctf_estimation.py @@ -1,4 +1,6 @@ +import os from logging import getLogger +from pathlib import Path from sqlmodel import Session, select @@ -12,12 +14,17 @@ 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 + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False @@ -37,7 +44,9 @@ def ctf_estimated(message: dict, murfey_db: Session) -> dict[str, bool]: )[session.instrument_name] if machine_config.smartem_api_url: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) update = MicrographUpdateRequest(status=MicrographStatus.CTF_COMPLETED) smartem_client._request( diff --git a/src/murfey/workflows/spa/flush_spa_preprocess.py b/src/murfey/workflows/spa/flush_spa_preprocess.py index 504a7940b..39d3547dc 100644 --- a/src/murfey/workflows/spa/flush_spa_preprocess.py +++ b/src/murfey/workflows/spa/flush_spa_preprocess.py @@ -1,4 +1,5 @@ import logging +import os from pathlib import Path from typing import Optional @@ -9,14 +10,19 @@ 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, ) + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False from murfey.server import _transport_object @@ -180,7 +186,9 @@ def register_grid_square( ) ) smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) gs_data = SmartEMGridSquareData( gridsquare_id=str(gsid), @@ -269,7 +277,9 @@ def register_foil_hole( )[murfey_session.instrument_name] if machine_config.smartem_api_url: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) fh_data = SmartEMFoilHoleData( id=str(foil_hole_params.name), diff --git a/src/murfey/workflows/spa/motion_correction.py b/src/murfey/workflows/spa/motion_correction.py index 0a50da167..8aa5c6a2d 100644 --- a/src/murfey/workflows/spa/motion_correction.py +++ b/src/murfey/workflows/spa/motion_correction.py @@ -1,4 +1,6 @@ +import os from logging import getLogger +from pathlib import Path from sqlmodel import Session, select @@ -12,12 +14,17 @@ 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 + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False @@ -37,7 +44,9 @@ def motion_corrected(message: dict, murfey_db: Session) -> dict[str, bool]: )[session.instrument_name] if machine_config.smartem_api_url: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) update = MicrographUpdateRequest( status=MicrographStatus.MOTION_CORRECTION_COMPLETED diff --git a/src/murfey/workflows/spa/picking.py b/src/murfey/workflows/spa/picking.py index f7c915d56..ec229fc30 100644 --- a/src/murfey/workflows/spa/picking.py +++ b/src/murfey/workflows/spa/picking.py @@ -1,4 +1,6 @@ +import os from logging import getLogger +from pathlib import Path from typing import List import numpy as np @@ -31,12 +33,17 @@ 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 + keycloak_client = KeycloakClient( + load_keycloak_config(Path(os.getenv("SMARTEM_KEYCLOAK_CONFIGURATION") or "")) + ) SMARTEM_ACTIVE = True except ImportError: + keycloak_client = None SMARTEM_ACTIVE = False @@ -390,7 +397,9 @@ def particles_picked(message: dict, murfey_db: Session) -> dict[str, bool]: )[session.instrument_name] if machine_config.smartem_api_url: smartem_client = SmartEMAPIClient( - base_url=machine_config.smartem_api_url, logger=logger + base_url=machine_config.smartem_api_url, + logger=logger, + keycloak_client=keycloak_client, ) update = MicrographUpdateRequest( status=MicrographStatus.PARTICLE_PICKING_COMPLETED