Wire in assets to the BrainAtlas and BrainAtlasRead (#207)

Author: mgeplf

app/schemas/brain_atlas.py

@@ -2,6 +2,7 @@
 
 from pydantic import BaseModel, ConfigDict
 
+from app.schemas.asset import AssetsMixin
 from app.schemas.base import CreationMixin, IdentifiableMixin, SpeciesRead
 
 
@@ -13,7 +14,7 @@ class BrainAtlasBase(BaseModel):
     species: SpeciesRead
 
 
-class BrainAtlasRead(BrainAtlasBase, CreationMixin, IdentifiableMixin):
+class BrainAtlasRead(BrainAtlasBase, CreationMixin, IdentifiableMixin, AssetsMixin):
     pass
 
 
@@ -27,5 +28,5 @@ class BrainAtlasRegionBase(BaseModel):
     brain_region_id: uuid.UUID
 
 
-class BrainAtlasRegionRead(BrainAtlasRegionBase, CreationMixin, IdentifiableMixin):
+class BrainAtlasRegionRead(BrainAtlasRegionBase, CreationMixin, IdentifiableMixin, AssetsMixin):
     pass

app/service/brain_atlas.py

@@ -1,5 +1,7 @@
 import uuid
 
+from sqlalchemy.orm import selectinload
+
 import app.queries.common
 from app.db.model import BrainAtlas, BrainAtlasRegion
 from app.dependencies.auth import UserContextDep
@@ -25,7 +27,7 @@ def read_many(
         facets=None,
         aliases=None,
         apply_filter_query_operations=None,
-        apply_data_query_operations=None,
+        apply_data_query_operations=lambda select: select.options(selectinload(BrainAtlas.assets)),
         pagination_request=pagination_request,
         response_schema_class=BrainAtlasRead,
         name_to_facet_query_params=None,
@@ -40,7 +42,7 @@ def read_one(user_context: UserContextDep, atlas_id: uuid.UUID, db: SessionDep)
         db_model_class=BrainAtlas,
         authorized_project_id=user_context.project_id,
         response_schema_class=BrainAtlasRead,
-        apply_operations=None,
+        apply_operations=lambda select: select.options(selectinload(BrainAtlas.assets)),
     )
 
 
@@ -62,7 +64,7 @@ def read_many_region(
         apply_filter_query_operations=lambda q: q.filter(
             BrainAtlasRegion.brain_atlas_id == atlas_id
         ),
-        apply_data_query_operations=None,
+        apply_data_query_operations=lambda s: s.options(selectinload(BrainAtlasRegion.assets)),
         pagination_request=pagination_request,
         response_schema_class=BrainAtlasRegionRead,
         name_to_facet_query_params=None,
@@ -79,5 +81,7 @@ def read_one_region(
         db_model_class=BrainAtlasRegion,
         authorized_project_id=user_context.project_id,
         response_schema_class=BrainAtlasRegionRead,
-        apply_operations=lambda q: q.filter(BrainAtlasRegion.brain_atlas_id == atlas_id),
+        apply_operations=lambda select: select.filter(
+            BrainAtlasRegion.brain_atlas_id == atlas_id
+        ).options(selectinload(BrainAtlasRegion.assets)),
     )

app/service/emodel.py

@@ -40,7 +40,6 @@ def _load(select: sa.Select[tuple[EModel]]):
         selectinload(EModel.contributions).joinedload(Contribution.role),
         joinedload(EModel.mtypes),
         joinedload(EModel.etypes),
-        selectinload(EModel.assets),
         selectinload(EModel.ion_channel_models).joinedload(IonChannelModel.species),
         selectinload(EModel.ion_channel_models).joinedload(IonChannelModel.strain),
         selectinload(EModel.ion_channel_models).joinedload(IonChannelModel.brain_region),

tests/routers/test_asset.py

@@ -1,12 +1,10 @@
 from unittest.mock import ANY
-from uuid import UUID
 
 import pytest
 
 from app.db.model import Asset, Entity
 from app.db.types import AssetLabel, AssetStatus, EntityType
 from app.errors import ApiErrorCode
-from app.routers.asset import EntityRoute
 from app.schemas.api import ErrorResponse
 from app.schemas.asset import AssetRead
 from app.utils.s3 import build_s3_path
@@ -18,6 +16,8 @@
     VIRTUAL_LAB_ID,
     add_db,
     create_reconstruction_morphology_id,
+    route,
+    upload_entity_asset,
 )
 
 DIFFERENT_ENTITY_TYPE = "experimental_bouton_density"
@@ -27,28 +27,6 @@
 FILE_EXAMPLE_SIZE = 31
 
 
-def _entity_type_to_route(entity_type: EntityType) -> EntityRoute:
-    return EntityRoute[entity_type.name]
-
-
-def _route(entity_type: EntityType) -> str:
-    return f"/{_entity_type_to_route(entity_type)}"
-
-
-def _upload_entity_asset(
-    client, entity_type: EntityType, entity_id: UUID, label: str | None = None
-):
-    with FILE_EXAMPLE_PATH.open("rb") as f:
-        files = {
-            # (filename, file (or bytes), content_type, headers)
-            "file": ("a/b/c.txt", f, "text/plain")
-        }
-        data = None
-        if label:
-            data = {"label": label}
-        return client.post(f"{_route(entity_type)}/{entity_id}/assets", files=files, data=data)
-
-
 def _get_expected_full_path(entity, path):
     return build_s3_path(
         vlab_id=VIRTUAL_LAB_ID,
@@ -73,6 +51,17 @@ def entity(client, species_id, strain_id, brain_region_id) -> Entity:
     return Entity(id=entity_id, type=entity_type)
 
 
+def _upload_entity_asset(client, entity_type, entity_id, label=None):
+    with FILE_EXAMPLE_PATH.open("rb") as f:
+        files = {
+            # (filename, file (or bytes), content_type, headers)
+            "file": ("a/b/c.txt", f, "text/plain")
+        }
+        return upload_entity_asset(
+            client=client, entity_type=entity_type, entity_id=entity_id, files=files, label=label
+        )
+
+
 @pytest.fixture
 def asset(client, entity) -> AssetRead:
     response = _upload_entity_asset(client, entity_type=entity.type, entity_id=entity.id)
@@ -178,7 +167,7 @@ def test_upload_entity_asset__label(monkeypatch, client, entity):
 
 
 def test_get_entity_asset(client, entity, asset):
-    response = client.get(f"{_route(entity.type)}/{entity.id}/assets/{asset.id}")
+    response = client.get(f"{route(entity.type)}/{entity.id}/assets/{asset.id}")
 
     assert response.status_code == 200, f"Failed to get asset: {response.text}"
     data = response.json()
@@ -197,20 +186,20 @@ def test_get_entity_asset(client, entity, asset):
     }
 
     # try to get an asset with non-existent entity id
-    response = client.get(f"{_route(entity.type)}/{MISSING_ID}/assets/{asset.id}")
+    response = client.get(f"{route(entity.type)}/{MISSING_ID}/assets/{asset.id}")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
     error = ErrorResponse.model_validate(response.json())
     assert error.error_code == ApiErrorCode.ENTITY_NOT_FOUND
 
     # try to get an asset with non-existent asset id
-    response = client.get(f"{_route(entity.type)}/{entity.id}/assets/{MISSING_ID}")
+    response = client.get(f"{route(entity.type)}/{entity.id}/assets/{MISSING_ID}")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
     error = ErrorResponse.model_validate(response.json())
     assert error.error_code == ApiErrorCode.ASSET_NOT_FOUND
 
 
 def test_get_entity_assets(client, entity, asset):
-    response = client.get(f"{_route(entity.type)}/{entity.id}/assets")
+    response = client.get(f"{route(entity.type)}/{entity.id}/assets")
 
     assert response.status_code == 200, f"Failed to get asset: {response.text}"
     data = response.json()["data"]
@@ -231,15 +220,15 @@ def test_get_entity_assets(client, entity, asset):
     ]
 
     # try to get assets with non-existent entity id
-    response = client.get(f"{_route(entity.type)}/{MISSING_ID}/assets")
+    response = client.get(f"{route(entity.type)}/{MISSING_ID}/assets")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
     error = ErrorResponse.model_validate(response.json())
     assert error.error_code == ApiErrorCode.ENTITY_NOT_FOUND
 
 
 def test_download_entity_asset(client, entity, asset):
     response = client.get(
-        f"{_route(entity.type)}/{entity.id}/assets/{asset.id}/download",
+        f"{route(entity.type)}/{entity.id}/assets/{asset.id}/download",
         follow_redirects=False,
     )
 
@@ -250,20 +239,20 @@ def test_download_entity_asset(client, entity, asset):
     assert expected_params.issubset(response.next_request.url.params)
 
     # try to download an asset with non-existent entity id
-    response = client.get(f"{_route(entity.type)}/{MISSING_ID}/assets/{asset.id}/download")
+    response = client.get(f"{route(entity.type)}/{MISSING_ID}/assets/{asset.id}/download")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
     error = ErrorResponse.model_validate(response.json())
     assert error.error_code == ApiErrorCode.ENTITY_NOT_FOUND
 
     # try to download an asset with non-existent asset id
-    response = client.get(f"{_route(entity.type)}/{entity.id}/assets/{MISSING_ID}/download")
+    response = client.get(f"{route(entity.type)}/{entity.id}/assets/{MISSING_ID}/download")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
     error = ErrorResponse.model_validate(response.json())
     assert error.error_code == ApiErrorCode.ASSET_NOT_FOUND
 
     # when downloading a single file asset_path should not be passed as a parameter
     response = client.get(
-        f"{_route(entity.type)}/{entity.id}/assets/{asset.id}/download",
+        f"{route(entity.type)}/{entity.id}/assets/{asset.id}/download",
         params={"asset_path": "foo"},
         follow_redirects=False,
     )
@@ -273,21 +262,21 @@ def test_download_entity_asset(client, entity, asset):
 
 
 def test_delete_entity_asset(client, entity, asset):
-    response = client.delete(f"{_route(entity.type)}/{entity.id}/assets/{asset.id}")
+    response = client.delete(f"{route(entity.type)}/{entity.id}/assets/{asset.id}")
     assert response.status_code == 200, f"Failed to delete asset: {response.text}"
     data = response.json()
     assert data == asset.model_copy(update={"status": AssetStatus.DELETED}).model_dump(mode="json")
 
     # try to delete again the same asset
-    response = client.delete(f"{_route(entity.type)}/{entity.id}/assets/{asset.id}")
+    response = client.delete(f"{route(entity.type)}/{entity.id}/assets/{asset.id}")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
 
     # try to delete an asset with non-existent entity id
-    response = client.delete(f"{_route(entity.type)}/{MISSING_ID}/assets/{asset.id}")
+    response = client.delete(f"{route(entity.type)}/{MISSING_ID}/assets/{asset.id}")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
 
     # try to delete an asset with non-existent asset id
-    response = client.delete(f"{_route(entity.type)}/{entity.id}/assets/{MISSING_ID}")
+    response = client.delete(f"{route(entity.type)}/{entity.id}/assets/{MISSING_ID}")
     assert response.status_code == 404, f"Unexpected result: {response.text}"
 
 
@@ -297,7 +286,7 @@ def test_upload_delete_upload_entity_asset(client, entity):
     data = response.json()
     asset0 = AssetRead.model_validate(data)
 
-    response = client.delete(f"{_route(entity.type)}/{entity.id}/assets/{asset0.id}")
+    response = client.delete(f"{route(entity.type)}/{entity.id}/assets/{asset0.id}")
     assert response.status_code == 200, f"Failed to delete asset: {response.text}"
 
     # upload the asset with the same path
@@ -307,7 +296,7 @@ def test_upload_delete_upload_entity_asset(client, entity):
     asset1 = AssetRead.model_validate(data)
 
     # test that the deleted assets are filtered out
-    response = client.get(f"{_route(entity.type)}/{entity.id}/assets")
+    response = client.get(f"{route(entity.type)}/{entity.id}/assets")
 
     assert response.status_code == 200, f"Failed to get assest: {response.text}"
     data = response.json()["data"]
@@ -320,15 +309,15 @@ def test_upload_delete_upload_entity_asset(client, entity):
 
 def test_download_directory_file(client, entity, asset_directory):
     response = client.get(
-        url=f"{_route(entity.type)}/{entity.id}/assets/{asset_directory.id}/download",
+        url=f"{route(entity.type)}/{entity.id}/assets/{asset_directory.id}/download",
         params={"asset_path": "file1.txt"},
         follow_redirects=False,
     )
     assert response.status_code == 307, f"Failed to download directory file: {response.text}"
 
     # asset_path is mandatory if the asset is a direcotory
     response = client.get(
-        url=f"{_route(entity.type)}/{entity.id}/assets/{asset_directory.id}/download",
+        url=f"{route(entity.type)}/{entity.id}/assets/{asset_directory.id}/download",
         follow_redirects=False,
     )
     assert response.status_code == 409, (