Cache API results (#23)

Co-authored-by: Reinder Vos de Wael <[email protected]>

api/.env.example

@@ -1,2 +1,3 @@
+ENVIRONMENT=development
 AZURE_STORAGE_BLOB_URL=https://ACCOUNT_NAME.blob.core.windows.net
 AZURE_ACCESS_KEY=SECRET_KEY

api/src/core/data.py

@@ -9,6 +9,8 @@
 from src.core import settings
 
 config = settings.get_settings()
+ENVIRONMENT = config.ENVIRONMENT
+DATA_DIR = config.DATA_DIR
 AZURE_STORAGE_BLOB_URL = config.AZURE_STORAGE_BLOB_URL
 AZURE_ACCESS_KEY = config.AZURE_ACCESS_KEY
 
@@ -63,6 +65,10 @@ def get_feature_data(species: str, side: str) -> np.ndarray:
     """
     logger.info("Getting feature file.")
     filename = f"{species}_{side}_gradient_10k_fs_lr.nii.gz"
+
+    if ENVIRONMENT == "development":
+        return nibabel.load(DATA_DIR / filename).get_fdata()  # type: ignore[attr-defined]
+
     with tempfile.NamedTemporaryFile(suffix=".nii.gz") as temp_file:
         download_file_from_blob(filename, temp_file.name)
         return nibabel.load(temp_file.name).get_fdata()  # type: ignore[attr-defined]
@@ -81,6 +87,9 @@ def get_surface_file(species: str, side: str) -> nibabel.GiftiImage:
     """
     logger.info("Getting surface file.")
     filename = f"{species}_{side}_inflated_10k_fs_lr.surf.gii"
+    if ENVIRONMENT == "development":
+        return nibabel.load(DATA_DIR / filename)  # type: ignore[return-value]
+
     with tempfile.NamedTemporaryFile(suffix=".surf.gii") as temp_file:
         download_file_from_blob(filename, temp_file.name)
         return nibabel.load(temp_file.name)  # type: ignore[return-value]

api/src/core/settings.py

@@ -1,6 +1,7 @@
 """Settings for the API."""
 import functools
 import logging
+import pathlib
 
 import pydantic
 
@@ -10,6 +11,11 @@ class Settings(pydantic.BaseSettings):  # type: ignore[valid-type, misc]
 
     LOGGER_NAME: str = pydantic.Field("Cross Species Mapper API", env="LOGGER_NAME")
 
+    ENVIRONMENT: str = pydantic.Field("development", env="ENVIRONMENT")
+    DATA_DIR: pathlib.Path = pydantic.Field(
+        pathlib.Path(__file__).parent.parent.parent / "data", env="DATA_DIR"
+    )
+
     AZURE_STORAGE_BLOB_URL: str = pydantic.Field(..., env="AZURE_STORAGE_BLOB_URL")
     AZURE_ACCESS_KEY: pydantic.SecretStr = pydantic.Field(..., env="AZURE_ACCESS_KEY")
 

api/src/core/utils.py

@@ -16,6 +16,25 @@
 logger = logging.getLogger(LOGGER_NAME)
 
 
+def add_cache_control(
+    response: fastapi.Response, expiry_in_minutes: int = 525600
+) -> fastapi.Response:
+    """Adds cache control headers to a response.
+
+    Args:
+        response: The response to add the headers to.
+        expiry_in_minutes: The number of minutes to set the expiry to.
+            Defaults to 525600 (1 year).
+
+    Returns:
+        The response with the headers added.
+
+    """
+    logger.info("Adding cache control headers.")
+    response.headers["Cache-Control"] = f"max-age={expiry_in_minutes * 60}"
+    return response
+
+
 class Surface:
     """Convenience class for surface handling."""
 

api/src/routers/features/views.py

@@ -6,7 +6,8 @@
 
 import fastapi
 from fastapi import status
-from src.core import settings
+
+from src.core import settings, utils
 from src.routers.features import controller, schemas
 
 router = fastapi.APIRouter(prefix="/features", tags=["features"])
@@ -21,6 +22,7 @@
     responses={status.HTTP_200_OK: {"model": List[schemas.FeatureSimilarity]}},
 )
 def get_feature_similarity(
+    response: fastapi.Response,
     seed_species: str = fastapi.Query(
         ..., example="human", description="The species to fetch the hemispheres for."
     ),
@@ -47,4 +49,5 @@ def get_feature_similarity(
         the seed vertex.
     """
     logger.info("Calling GET /surfaces/similarity endpoint.")
+    response = utils.add_cache_control(response)
     return controller.get_cross_species_features(seed_species, seed_side, seed_vertex)

api/src/routers/surfaces/views.py

@@ -3,7 +3,8 @@
 
 import fastapi
 from fastapi import status
-from src.core import settings
+
+from src.core import settings, utils
 from src.routers.surfaces import controller, schemas
 
 config = settings.get_settings()
@@ -16,6 +17,7 @@
 
 @router.get("/hemispheres", responses={status.HTTP_200_OK: {"model": schemas.Surface}})
 def get_hemispheres(
+    response: fastapi.Response,
     species: str = fastapi.Query(
         ..., example="human", description="The species to fetch the hemispheres for."
     ),
@@ -35,4 +37,5 @@ def get_hemispheres(
         A Pydantic BaseClass containing the surface.
     """
     logger.info("Calling GET /surfaces/hemispheres endpoint.")
+    response = utils.add_cache_control(response)
     return controller.get_hemispheres(species, side)