Merge remote-tracking branch 'origin/main' into simplified-brain-regions

Author: mgeplf

alembic/versions/20250509_121806_bcc13861444b_base_models.py

@@ -35,6 +35,7 @@
     AgentType,
     AgePeriod,
     AnnotationBodyType,
+    AssetLabel,
     AssetStatus,
     ElectricalRecordingOrigin,
     ElectricalRecordingStimulusShape,
@@ -852,6 +853,7 @@ class Asset(Identifiable):
     size: Mapped[BIGINT]
     sha256_digest: Mapped[bytes | None] = mapped_column(LargeBinary(32))
     meta: Mapped[JSON_DICT]  # not used yet. can be useful?
+    label: Mapped[AssetLabel | None]
     entity_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("entity.id"), index=True)
 
     # partial unique index

alembic/versions/20250509_121821_f81f202805ce_triggers.py

@@ -176,3 +176,18 @@ class StructuralDomain(StrEnum):
     apical_dendrite = auto()
     basal_dendrite = auto()
     axon = auto()
+
+
+class AssetLabel(StrEnum):
+    neurolucida = auto()
+    swc = auto()
+    hdf5 = auto()
+
+
+ALLOWED_ASSET_LABELS_PER_ENTITY = {
+    EntityType.reconstruction_morphology: {
+        AssetLabel.neurolucida,
+        AssetLabel.swc,
+        AssetLabel.hdf5,
+    }
+}

app/db/model.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 from psycopg2.errors import ForeignKeyViolation, InsufficientPrivilege, UniqueViolation
+from pydantic import ValidationError
 from sqlalchemy.exc import IntegrityError, NoResultFound, ProgrammingError
 
 from app.utils.enum import UpperStrEnum
@@ -43,6 +44,7 @@ class ApiErrorCode(UpperStrEnum):
     ASSET_MISSING_PATH = auto()
     ASSET_INVALID_PATH = auto()
     ASSET_NOT_A_DIRECTORY = auto()
+    ASSET_INVALID_SCHEMA = auto()
     ION_NAME_NOT_FOUND = auto()
 
 
@@ -138,3 +140,19 @@ def ensure_authorized_references(
                 message=error_message, error_code=error_code, http_status_code=HTTPStatus.FORBIDDEN
             ) from err
         raise
+
+
+@contextmanager
+def ensure_valid_schema(
+    error_message: str, error_code: ApiErrorCode = ApiErrorCode.INVALID_REQUEST
+) -> Iterator[None]:
+    """Context manager that raises ApiError when a schema validation error is raised."""
+    try:
+        yield
+    except ValidationError as err:
+        raise ApiError(
+            message=error_message,
+            error_code=error_code,
+            http_status_code=HTTPStatus.UNPROCESSABLE_ENTITY,
+            details=[e["msg"] for e in err.errors()],
+        ) from err

app/db/types.py

@@ -65,6 +65,7 @@ def create_entity_asset(self, entity_id: uuid.UUID, asset: AssetCreate) -> Asset
                 size=asset.size,
                 sha256_digest=sha256_digest,
                 meta=asset.meta,
+                label=asset.label,
             )
             .returning(Asset)
         )

app/errors.py

@@ -10,7 +10,7 @@
 from starlette.responses import RedirectResponse
 
 from app.config import settings
-from app.db.types import EntityType
+from app.db.types import AssetLabel, EntityType
 from app.dependencies.auth import UserContextDep, UserContextWithProjectIdDep
 from app.dependencies.db import RepoGroupDep
 from app.dependencies.s3 import S3ClientDep
@@ -92,6 +92,7 @@ def upload_entity_asset(
     entity_id: uuid.UUID,
     file: UploadFile,
     meta: Annotated[dict | None, Form()] = None,
+    label: Annotated[AssetLabel | None, Form()] = None,
 ) -> AssetRead:
     """Upload an asset to be associated with the specified entity.
 
@@ -123,6 +124,7 @@ def upload_entity_asset(
         size=file.size,
         sha256_digest=sha256_digest,
         meta=meta,
+        label=label,
     )
     if not upload_to_s3(
         s3_client,

app/repository/asset.py

@@ -1,8 +1,8 @@
 import uuid
 
-from pydantic import BaseModel, ConfigDict, field_validator
+from pydantic import BaseModel, ConfigDict, field_validator, model_validator
 
-from app.db.types import AssetStatus
+from app.db.types import ALLOWED_ASSET_LABELS_PER_ENTITY, AssetLabel, AssetStatus, EntityType
 
 
 class AssetBase(BaseModel):
@@ -16,6 +16,7 @@ class AssetBase(BaseModel):
     size: int
     sha256_digest: str | None
     meta: dict
+    label: AssetLabel | None = None
 
     @field_validator("sha256_digest", mode="before")
     @classmethod
@@ -35,6 +36,29 @@ class AssetRead(AssetBase):
 class AssetCreate(AssetBase):
     """Asset model for creation."""
 
+    entity_type: EntityType
+
+    @model_validator(mode="after")
+    def ensure_entity_type_label_consistency(self):
+        """Asset label must be within the allowed labels for the entity_type."""
+        if not self.label:
+            return self
+
+        allowed_asset_labels = ALLOWED_ASSET_LABELS_PER_ENTITY.get(self.entity_type, None)
+
+        if allowed_asset_labels is None:
+            msg = f"There are no allowed asset labels defined for '{self.entity_type}'"
+            raise ValueError(msg)
+
+        if self.label not in allowed_asset_labels:
+            msg = (
+                f"Asset label '{self.label}' is not allowed for entity type '{self.entity_type}'. "
+                f"Allowed asset labels: {sorted(label.value for label in allowed_asset_labels)}"
+            )
+            raise ValueError(msg)
+
+        return self
+
 
 class AssetsMixin(BaseModel):
     assets: list[AssetRead] | None

app/routers/asset.py

@@ -1,7 +1,7 @@
 import uuid
 
-from app.db.types import AssetStatus, EntityType
-from app.errors import ApiErrorCode, ensure_result, ensure_uniqueness
+from app.db.types import AssetLabel, AssetStatus, EntityType
+from app.errors import ApiErrorCode, ensure_result, ensure_uniqueness, ensure_valid_schema
 from app.repository.group import RepositoryGroup
 from app.schemas.asset import AssetCreate, AssetRead
 from app.schemas.auth import UserContext, UserContextWithProjectId
@@ -59,6 +59,7 @@ def create_entity_asset(
     size: int,
     sha256_digest: str | None,
     meta: dict | None,
+    label: AssetLabel | None,
 ) -> AssetRead:
     """Create an asset for an entity."""
     entity = entity_service.get_writable_entity(
@@ -75,15 +76,21 @@ def create_entity_asset(
         filename=filename,
         is_public=entity.authorized_public,
     )
-    asset_create = AssetCreate(
-        path=filename,
-        full_path=full_path,
-        is_directory=False,
-        content_type=content_type,
-        size=size,
-        sha256_digest=sha256_digest,
-        meta=meta or {},
-    )
+
+    with ensure_valid_schema(
+        "Asset schema is invalid", error_code=ApiErrorCode.ASSET_INVALID_SCHEMA
+    ):
+        asset_create = AssetCreate(
+            path=filename,
+            full_path=full_path,
+            is_directory=False,
+            content_type=content_type,
+            size=size,
+            sha256_digest=sha256_digest,
+            meta=meta or {},
+            label=label,
+            entity_type=entity_type,
+        )
     with ensure_uniqueness(
         f"Asset with path {asset_create.path!r} already exists",
         error_code=ApiErrorCode.ASSET_DUPLICATED,

app/schemas/asset.py

@@ -4,7 +4,7 @@
 import pytest
 
 from app.db.model import Asset, Entity
-from app.db.types import AssetStatus, EntityType
+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
@@ -35,13 +35,18 @@ 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):
+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")
         }
-        return client.post(f"{_route(entity_type)}/{entity_id}/assets", files=files)
+        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):
@@ -95,7 +100,9 @@ def asset_directory(db, entity) -> AssetRead:
 
 
 def test_upload_entity_asset(client, entity):
-    response = _upload_entity_asset(client, entity_type=entity.type, entity_id=entity.id)
+    response = _upload_entity_asset(
+        client, entity_type=entity.type, entity_id=entity.id, label="neurolucida"
+    )
     assert response.status_code == 201, f"Failed to create asset: {response.text}"
     data = response.json()
 
@@ -110,6 +117,7 @@ def test_upload_entity_asset(client, entity):
         "sha256_digest": FILE_EXAMPLE_DIGEST,
         "meta": {},
         "status": "created",
+        "label": "neurolucida",
     }
 
     # try to upload again the same file with the same path
@@ -133,6 +141,42 @@ def test_upload_entity_asset(client, entity):
     assert error.error_code == ApiErrorCode.ENTITY_NOT_FOUND
 
 
+def test_upload_entity_asset__label(monkeypatch, client, entity):
+    response = _upload_entity_asset(
+        client, entity_type=entity.type, entity_id=entity.id, label="foo"
+    )
+    assert response.status_code == 422, "Assel label was not rejected as not present in AssetLabel."
+
+    monkeypatch.setattr("app.schemas.asset.ALLOWED_ASSET_LABELS_PER_ENTITY", {})
+
+    response = _upload_entity_asset(
+        client, entity_type=entity.type, entity_id=entity.id, label=AssetLabel.hdf5
+    )
+    assert response.status_code == 422
+    assert response.json() == {
+        "error_code": "ASSET_INVALID_SCHEMA",
+        "message": "Asset schema is invalid",
+        "details": [f"Value error, There are no allowed asset labels defined for '{entity.type}'"],
+    }
+
+    required = {EntityType.reconstruction_morphology: {AssetLabel.swc}}
+
+    monkeypatch.setattr("app.schemas.asset.ALLOWED_ASSET_LABELS_PER_ENTITY", required)
+
+    response = _upload_entity_asset(
+        client, entity_type=entity.type, entity_id=entity.id, label=AssetLabel.hdf5
+    )
+    assert response.status_code == 422
+    assert response.json() == {
+        "error_code": "ASSET_INVALID_SCHEMA",
+        "message": "Asset schema is invalid",
+        "details": [
+            f"Value error, Asset label '{AssetLabel.hdf5}' is not allowed for "
+            f"entity type '{entity.type}'. Allowed asset labels: ['{AssetLabel.swc}']"
+        ],
+    }
+
+
 def test_get_entity_asset(client, entity, asset):
     response = client.get(f"{_route(entity.type)}/{entity.id}/assets/{asset.id}")
 
@@ -149,6 +193,7 @@ def test_get_entity_asset(client, entity, asset):
         "sha256_digest": FILE_EXAMPLE_DIGEST,
         "meta": {},
         "status": "created",
+        "label": None,
     }
 
     # try to get an asset with non-existent entity id
@@ -181,6 +226,7 @@ def test_get_entity_assets(client, entity, asset):
             "sha256_digest": FILE_EXAMPLE_DIGEST,
             "meta": {},
             "status": "created",
+            "label": None,
         }
     ]