Add auth for smartem

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={

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

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

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,15 +22,20 @@
 
 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,
         MicrographData as SmartEMMicrographData,
         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(

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(

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(

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),

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

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