Added FastAPI endpoints to instrument and backend servers to request for general file downloads from different upstream instruments

Author: tieneupin

src/murfey/instrument_server/api.py

@@ -469,13 +469,83 @@ def upload_gain_reference(
     return {"success": True}
 
 
-class UpstreamTiffInfo(BaseModel):
+class UpstreamFileDownloadInfo(BaseModel):
+    download_dir: Path
+    upstream_instrument: str
+    upstream_visit_path: Path
+
+
+@router.post("/visits/{visit_name}/sessions/{session_id}/upstream_file_data_request")
+def gather_upstream_files(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_file_download: UpstreamFileDownloadInfo,
+):
+    """
+    Instrument server endpoint that will query the backend for files in the chosen
+    visit directory
+    """
+    # Check for forbidden characters
+    if any(c in visit_name for c in ("/", "\\", ":", ";")):
+        logger.error(f"Forbidden characters are present in the visit name {visit_name}")
+        return {
+            "succss": False,
+            "detail": "Forbidden characters present in visit name",
+        }
+    # Get the list of files to download
+    murfey_url = urlparse(_get_murfey_url(), allow_fragments=False)
+    sanitised_visit_name = sanitise_nonpath(visit_name)
+    url_path = url_path_for(
+        "session_control.correlative_router",
+        "gather_upstream_files",
+        session_id=session_id,
+        visit_name=sanitised_visit_name,
+    )
+    upstream_files: list[str] = requests.get(
+        f"{murfey_url.geturl()}{url_path}",
+        headers={"Authorization": f"Bearer {tokens[session_id]}"},
+        json={
+            "upstream_instrument": upstream_file_download.upstream_instrument,
+            "upstream_visit_path": str(upstream_file_download.upstream_visit_path),
+        },
+    ).json()
+
+    # Make the download directory and download gathered files
+    upstream_file_download.download_dir.mkdir(exist_ok=True)
+    for upstream_file in upstream_files:
+        url_path = url_path_for(
+            "session_control.correlative_router",
+            "get_upstream_file",
+            session_id=session_id,
+            visit_name=sanitised_visit_name,
+            upstream_file_path=upstream_file,
+        )
+        file_data = requests.get(
+            f"{murfey_url.geturl()}{url_path}",
+            headers={"Authorization": f"Bearer {tokens[session_id]}"},
+            stream=True,
+        )
+        upstream_file_relative_path = Path(upstream_file).relative_to(
+            upstream_file_download.upstream_visit_path
+        )
+        save_file = upstream_file_download.download_dir / upstream_file_relative_path
+        save_file.parent.mkdir(parents=True, exist_ok=True)
+        with open(save_file, "wb") as f:
+            for chunk in file_data.iter_content(chunk_size=32 * 1024**2):
+                f.write(chunk)
+        logger.info(f"Saved file to {str(save_file)!r}")
+    return {"success": True}
+
+
+class UpstreamTiffDownloadInfo(BaseModel):
     download_dir: Path
 
 
 @router.post("/visits/{visit_name}/sessions/{session_id}/upstream_tiff_data_request")
 def gather_upstream_tiffs(
-    visit_name: str, session_id: MurfeySessionID, upstream_tiff_info: UpstreamTiffInfo
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_tiff_info: UpstreamTiffDownloadInfo,
 ):
     sanitised_visit_name = sanitise_nonpath(visit_name)
     assert not any(c in visit_name for c in ("/", "\\", ":", ";"))

src/murfey/server/api/instrument.py

@@ -396,6 +396,69 @@ async def request_upstream_tiff_data_download(
     return data
 
 
+class UpstreamFileRequestInfo(BaseModel):
+    upstream_instrument: str
+    upstream_visit_path: Path
+
+
+@router.post("/visits/{visit_name}/sessions/{session_id}/upstream_file_data_request")
+async def request_upstream_file_data_download(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_file_request: UpstreamFileRequestInfo,
+    db=murfey_db,
+):
+    """
+    Forwards a request to the instrument server to trigger a file download request.
+    """
+    # Load the current instrument's machine config
+    instrument_name = (
+        db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
+    )
+    machine_config = get_machine_config(instrument_name=instrument_name)[
+        instrument_name
+    ]
+
+    # Log and return errors if download directory or URL weren't specified
+    if not machine_config.upstream_data_download_directory:
+        log.error("No download directory was configured for this instrument")
+        return {
+            "success": False,
+            "detail": "No download directory configured",
+        }
+    if not machine_config.instrument_server_url:
+        log.error("Couldn't find instrument server URL to post request to")
+        return {
+            "success": False,
+            "detail": "No instrument server URL",
+        }
+
+    # Forward the download request
+    download_dir = str(
+        machine_config.upstream_data_download_directory / secure_filename(visit_name)
+    )
+    async with aiohttp.ClientSession() as clientsession:
+        url_path = url_path_for(
+            "api.router",
+            "gather_upstream_files",
+            visit_name=secure_filename(visit_name),
+            session_id=session_id,
+        )
+        async with clientsession.post(
+            f"{machine_config.instrument_server_url}{url_path}",
+            headers={
+                "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}"
+            },
+            json={
+                "download_dir": download_dir,
+                "upstream_instrument": upstream_file_request.upstream_instrument,
+                "upstream_visit_path": str(upstream_file_request.upstream_visit_path),
+            },
+        ) as resp:
+            data = await resp.json()
+    return data
+
+
 class RsyncerSource(BaseModel):
     source: str
 

src/murfey/server/api/session_control.py

@@ -18,13 +18,15 @@
 )
 from murfey.server.api.shared import (
     find_upstream_visits as _find_upstream_visits,
+    gather_upstream_files as _gather_upstream_files,
     gather_upstream_tiffs as _gather_upstream_tiffs,
     get_foil_hole as _get_foil_hole,
     get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square,
     get_grid_squares as _get_grid_squares,
     get_grid_squares_from_dcg as _get_grid_squares_from_dcg,
     get_machine_config_for_instrument,
     get_tiff_file as _get_tiff_file,
+    get_upstream_file as _get_upstream_file,
     remove_session_by_id,
 )
 from murfey.server.ispyb import DB as ispyb_db, get_all_ongoing_visits
@@ -422,6 +424,43 @@ async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db):
     return _find_upstream_visits(session_id=session_id, db=db)
 
 
+class UpstreamFileGatheringInfo(BaseModel):
+    upstream_instrument: str
+    upstream_visit_path: Path
+
+
+@correlative_router.get(
+    "/visits/{visit_name}/sessions/{session_id}/upstream_file_paths"
+)
+async def gather_upstream_files(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_file_gathering: UpstreamFileGatheringInfo,
+    db=murfey_db,
+):
+    return _gather_upstream_files(
+        session_id=session_id,
+        upstream_instrument=upstream_file_gathering.upstream_instrument,
+        upstream_visit_path=upstream_file_gathering.upstream_visit_path,
+        db=db,
+    )
+
+
+@correlative_router.get(
+    "/visits/{visit_name}/sessions/{session_id}/upstream_file/{upstream_file_path:path}"
+)
+async def get_upstream_file(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_file_path: str,
+    db=murfey_db,
+):
+    upstream_file = _get_upstream_file(upstream_file_path)
+    return (
+        FileResponse(path=upstream_file) if upstream_file is not None else upstream_file
+    )
+
+
 @correlative_router.get(
     "/visits/{visit_name}/sessions/{session_id}/upstream_tiff_paths"
 )
@@ -436,4 +475,4 @@ async def get_tiff_file(visit_name: str, session_id: int, tiff_path: str, db=mur
     tiff_file = _get_tiff_file(
         visit_name=visit_name, session_id=session_id, tiff_path=tiff_path, db=db
     )
-    return FileResponse(path=tiff_file) if isinstance(tiff_file, Path) else tiff_file
+    return FileResponse(path=tiff_file) if tiff_file is not None else tiff_file

src/murfey/server/api/session_info.py

@@ -18,13 +18,15 @@
 )
 from murfey.server.api.shared import (
     find_upstream_visits as _find_upstream_visits,
+    gather_upstream_files as _gather_upstream_files,
     gather_upstream_tiffs as _gather_upstream_tiffs,
     get_foil_hole as _get_foil_hole,
     get_foil_holes_from_grid_square as _get_foil_holes_from_grid_square,
     get_grid_squares as _get_grid_squares,
     get_grid_squares_from_dcg as _get_grid_squares_from_dcg,
     get_machine_config_for_instrument,
     get_tiff_file as _get_tiff_file,
+    get_upstream_file as _get_upstream_file,
     remove_session_by_id,
 )
 from murfey.server.ispyb import DB as ispyb_db, get_all_ongoing_visits
@@ -416,6 +418,43 @@ async def find_upstream_visits(session_id: MurfeySessionID, db=murfey_db):
     return _find_upstream_visits(session_id=session_id, db=db)
 
 
+class UpstreamFileGatheringInfo(BaseModel):
+    upstream_instrument: str
+    upstream_visit_path: Path
+
+
+@correlative_router.get(
+    "/visits/{visit_name}/sessions/{session_id}/upstream_file_paths"
+)
+async def gather_upstream_files(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_file_gathering: UpstreamFileGatheringInfo,
+    db=murfey_db,
+):
+    return _gather_upstream_files(
+        session_id=session_id,
+        upstream_instrument=upstream_file_gathering.upstream_instrument,
+        upstream_visit_path=upstream_file_gathering.upstream_visit_path,
+        db=db,
+    )
+
+
+@correlative_router.get(
+    "/visits/{visit_name}/sessions/{session_id}/upstream_file/{upstream_file_path:path}"
+)
+async def get_upstream_file(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    upstream_file_path: Path,
+    db=murfey_db,
+):
+    upstream_file = _get_upstream_file(upstream_file_path)
+    return (
+        FileResponse(path=upstream_file) if upstream_file is not None else upstream_file
+    )
+
+
 @correlative_router.get(
     "/visits/{visit_name}/sessions/{session_id}/upstream_tiff_paths"
 )

src/murfey/server/api/shared.py

@@ -168,13 +168,62 @@ def find_upstream_visits(session_id: int, db: SQLModelSession):
         # Looks for visit name in file path
         current_upstream_visits = {}
         for visit_path in Path(upstream_data_dir).glob(f"{visit_name.split('-')[0]}-*"):
-            current_upstream_visits[visit_path.name] = (
-                visit_path / machine_config.processed_directory_name
-            )
+            current_upstream_visits[visit_path.name] = visit_path
         upstream_visits[upstream_instrument] = current_upstream_visits
     return upstream_visits
 
 
+def gather_upstream_files(
+    session_id: int,
+    upstream_instrument: str,
+    upstream_visit_path: Path,
+    db: SQLModelSession,
+):
+    """
+    Searches the specified upstream instrument for files based on the search strings
+    set in the MachineConfig and retursn them as a list of file paths.
+    """
+    # Load the current instrument's machine config
+    murfey_session = db.exec(
+        select(MurfeySession).where(MurfeySession.id == session_id)
+    ).one()
+    instrument_name = murfey_session.instrument_name
+    machine_config = get_machine_config(instrument_name=instrument_name)[
+        instrument_name
+    ]
+
+    # Search for files using the configured strings for that upstream instrument
+    file_list: list[Path] = []
+    logger.info(f"Searching for files in {upstream_visit_path}")
+    if (
+        machine_config.upstream_data_search_strings.get(upstream_instrument, None)
+        is not None
+    ):
+        for search_string in machine_config.upstream_data_search_strings[
+            upstream_instrument
+        ]:
+            logger.info(f"Using search string {search_string}")
+            for file in upstream_visit_path.glob(search_string):
+                if file.is_file():
+                    file_list.append(file)
+        logger.info(
+            f"Found {len(file_list)} files for download from {upstream_instrument}"
+        )
+    else:
+        logger.warning(
+            f"Upstream file searching has not been configured for {upstream_instrument} on {instrument_name}"
+        )
+    return file_list
+
+
+def get_upstream_file(file_path: str | Path):
+    file_path = Path(file_path) if isinstance(file_path, str) else file_path
+    if file_path.exists() and file_path.is_file():
+        return file_path
+    logger.warning(f"Requested file {str(file_path)!r} was not found")
+    return None
+
+
 def get_upstream_tiff_dirs(visit_name: str, instrument_name: str) -> List[Path]:
     tiff_dirs = []
     machine_config = get_machine_config(instrument_name=instrument_name)[