Add support for admin endpoints (#128)

Author: eleftherioszisis

src/entitysdk/client.py

@@ -38,6 +38,7 @@
 )
 from entitysdk.utils.asset import filter_assets
 
+TEntity = TypeVar("TEntity", bound=Entity)
 TIdentifiable = TypeVar("TIdentifiable", bound=Identifiable)
 
 
@@ -112,13 +113,15 @@ def get_entity(
         *,
         entity_type: type[TIdentifiable],
         project_context: ProjectContext | None = None,
+        admin: bool = False,
     ) -> TIdentifiable:
         """Get entity from resource id.
 
         Args:
             entity_id: Resource id of the entity.
             entity_type: Type of the entity.
             project_context: Optional project context.
+            admin: Whether to use the admin endpoint or not.
 
         Returns:
             entity_type instantiated by deserializing the response.
@@ -127,8 +130,11 @@ def get_entity(
             api_url=self.api_url,
             entity_type=entity_type,
             entity_id=entity_id,
+            admin=admin,
+        )
+        context = (
+            self._optional_user_context(override_context=project_context) if not admin else None
         )
-        context = self._optional_user_context(override_context=project_context)
         return core.get_entity(
             url=url,
             entity_type=entity_type,
@@ -208,13 +214,36 @@ def register_entity(
             token=self._token_manager.get_token(),
         )
 
+    def get_entity_assets(
+        self,
+        entity_id: ID,
+        *,
+        entity_type: type[TEntity],
+        project_context: ProjectContext | None = None,
+        admin: bool = False,
+    ) -> IteratorResult[Asset]:
+        """Get all assets of an entity."""
+        context = (
+            self._optional_user_context(override_context=project_context) if not admin else None
+        )
+        return core.get_entity_assets(
+            api_url=self.api_url,
+            entity_id=entity_id,
+            entity_type=entity_type,
+            project_context=context,
+            http_client=self._http_client,
+            token=self._token_manager.get_token(),
+            admin=admin,
+        )
+
     def update_entity(
         self,
         entity_id: ID,
         entity_type: type[TIdentifiable],
         attrs_or_entity: dict | Identifiable,
         *,
         project_context: ProjectContext | None = None,
+        admin: bool = False,
     ) -> TIdentifiable:
         """Update an entity.
 
@@ -223,13 +252,17 @@ def update_entity(
             entity_type: Type of the entity.
             attrs_or_entity: Attributes or entity to update.
             project_context: Optional project context.
+            admin: whether to use the admin endpoint or not.
         """
         url = route.get_entities_endpoint(
             api_url=self.api_url,
             entity_type=entity_type,
             entity_id=entity_id,
+            admin=admin,
+        )
+        context = (
+            self._optional_user_context(override_context=project_context) if not admin else None
         )
-        context = self._required_user_context(override_context=project_context)
         return core.update_entity(
             url=url,
             project_context=context,
@@ -239,6 +272,27 @@ def update_entity(
             token=self._token_manager.get_token(),
         )
 
+    def delete_entity(
+        self,
+        entity_id: ID,
+        entity_type: type[Identifiable],
+        *,
+        admin: bool = False,
+    ) -> None:
+        """Delete an entity."""
+        url = route.get_entities_endpoint(
+            api_url=self.api_url,
+            entity_type=entity_type,
+            entity_id=entity_id,
+            admin=admin,
+        )
+        core.delete_entity(
+            url=url,
+            entity_type=entity_type,
+            http_client=self._http_client,
+            token=self._token_manager.get_token(),
+        )
+
     def upload_file(
         self,
         *,
@@ -617,15 +671,19 @@ def delete_asset(
         entity_type: type[Identifiable],
         asset_id: ID,
         project_context: ProjectContext | None = None,
+        admin: bool = False,
     ) -> Asset:
         """Delete an entity's asset."""
         url = route.get_assets_endpoint(
             api_url=self.api_url,
             entity_type=entity_type,
             entity_id=entity_id,
             asset_id=asset_id,
+            admin=admin,
+        )
+        context = (
+            self._required_user_context(override_context=project_context) if not admin else None
         )
-        context = self._required_user_context(override_context=project_context)
         return core.delete_asset(
             url=url,
             project_context=context,

src/entitysdk/core.py

@@ -21,7 +21,7 @@
 from entitysdk.models.core import Identifiable
 from entitysdk.models.entity import Entity
 from entitysdk.result import IteratorResult
-from entitysdk.route import get_entity_derivations_endpoint
+from entitysdk.route import get_assets_endpoint, get_entity_derivations_endpoint
 from entitysdk.types import ID, AssetLabel, DerivationType
 from entitysdk.util import make_db_api_request, stream_paginated_request
 
@@ -74,13 +74,15 @@ def get_entity(
     entity_type: type[TIdentifiable],
     project_context: ProjectContext | None = None,
     token: str,
+    options: dict | None = None,
     http_client: httpx.Client | None = None,
 ) -> TIdentifiable:
     """Instantiate entity with model ``entity_type`` from resource id."""
     response = make_db_api_request(
         url=url,
         method="GET",
         json=None,
+        parameters=options,
         project_context=project_context,
         token=token,
         http_client=http_client,
@@ -121,6 +123,36 @@ def get_entity_derivations(
     )
 
 
+def get_entity_assets(
+    *,
+    api_url: str,
+    entity_id: ID,
+    entity_type: type[Entity],
+    project_context: ProjectContext | None,
+    token: str,
+    http_client: httpx.Client | None = None,
+    admin: bool = False,
+):
+    """Get all assets of an entity."""
+    url = get_assets_endpoint(
+        api_url=api_url,
+        entity_type=entity_type,
+        entity_id=entity_id,
+        asset_id=None,
+        admin=admin,
+    )
+    response = make_db_api_request(
+        url=url,
+        method="GET",
+        project_context=project_context,
+        token=token,
+        http_client=http_client,
+    )
+    return IteratorResult(
+        serdes.deserialize_model(json_data, Asset) for json_data in response.json()["data"]
+    )
+
+
 def register_entity(
     url: str,
     *,
@@ -148,7 +180,7 @@ def update_entity(
     *,
     entity_type: type[TIdentifiable],
     attrs_or_entity: dict | Identifiable,
-    project_context: ProjectContext,
+    project_context: ProjectContext | None,
     token: str,
     http_client: httpx.Client | None = None,
 ) -> TIdentifiable:
@@ -172,6 +204,22 @@ def update_entity(
     return serdes.deserialize_model(json_data, entity_type)
 
 
+def delete_entity(
+    url: str,
+    *,
+    entity_type: type[Identifiable],
+    token: str,
+    http_client: httpx.Client | None = None,
+):
+    """Delete entity."""
+    make_db_api_request(
+        url=url,
+        method="DELETE",
+        token=token,
+        http_client=http_client,
+    )
+
+
 def upload_asset_file(
     url: str,
     *,
@@ -372,7 +420,7 @@ def download_asset_content(
 def delete_asset(
     url: str,
     *,
-    project_context: ProjectContext,
+    project_context: ProjectContext | None,
     token: str,
     http_client: httpx.Client | None = None,
 ) -> Asset:

src/entitysdk/models/__init__.py

@@ -23,6 +23,7 @@
 from entitysdk.models.ion_channel_model import IonChannelModel, NeuronBlock, UseIon
 from entitysdk.models.ion_channel_recording import IonChannelRecording
 from entitysdk.models.license import License
+from entitysdk.models.measurement_annotation import MeasurementAnnotation
 from entitysdk.models.memodel import MEModel
 from entitysdk.models.memodelcalibrationresult import MEModelCalibrationResult
 from entitysdk.models.mtype import MTypeClass
@@ -63,6 +64,7 @@
     "IonChannelModel",
     "IonChannelRecording",
     "License",
+    "MeasurementAnnotation",
     "MEModel",
     "MEModelCalibrationResult",
     "MTypeClass",

src/entitysdk/models/measurement_annotation.py

@@ -0,0 +1,35 @@
+"""Measurement annotation."""
+
+from entitysdk.models.base import BaseModel
+from entitysdk.models.core import Identifiable
+from entitysdk.types import (
+    ID,
+    MeasurableEntity,
+    MeasurementStatistic,
+    MeasurementUnit,
+    StructuralDomain,
+)
+
+
+class MeasurementItem(BaseModel):
+    """Measurement item."""
+
+    name: MeasurementStatistic | None
+    unit: MeasurementUnit | None
+    value: float | None
+
+
+class MeasurementKind(BaseModel):
+    """Measurement kind."""
+
+    structural_domain: StructuralDomain
+    measurement_items: list[MeasurementItem]
+    pref_label: str
+
+
+class MeasurementAnnotation(Identifiable):
+    """Measurement annotation."""
+
+    entity_id: ID
+    entity_type: MeasurableEntity
+    measurement_kinds: list[MeasurementKind]

src/entitysdk/route.py

@@ -29,6 +29,7 @@
     "IonChannelModel": "ion-channel-model",
     "IonChannelRecording": "ion-channel-recording",
     "License": "license",
+    "MeasurementAnnotation": "measurement-annotation",
     "MEModel": "memodel",
     "MEModelCalibrationResult": "memodel-calibration-result",
     "MTypeClassification": "mtype-classification",
@@ -79,9 +80,14 @@ def get_entities_endpoint(
     api_url: str,
     entity_type: type[Identifiable],
     entity_id: str | ID | None = None,
+    admin: bool = False,
 ) -> str:
     """Get the API endpoint for an entity type."""
     route_name = get_route_name(entity_type)
+
+    if admin:
+        route_name = f"admin/{route_name}"
+
     endpoint = route_name if entity_id is None else f"{route_name}/{entity_id}"
     return f"{api_url}/{endpoint}"
 
@@ -103,6 +109,7 @@ def get_assets_endpoint(
     entity_type: type[Identifiable],
     entity_id: str | ID,
     asset_id: str | ID | None = None,
+    admin: bool = False,
 ) -> str:
     """Return the endpoint for the assets of an entity.
 
@@ -111,10 +118,13 @@ def get_assets_endpoint(
         entity_type: The type of the entity.
         entity_id: The ID of the entity.
         asset_id: The ID of the asset.
+        admin: If true route is prefixed by admin, e.g. /admin/entity
 
     Returns:
         The endpoint for the assets of an entity.
     """
-    base_url = get_entities_endpoint(api_url=api_url, entity_type=entity_type, entity_id=entity_id)
+    base_url = get_entities_endpoint(
+        api_url=api_url, entity_type=entity_type, entity_id=entity_id, admin=admin
+    )
     asset_path = "assets" if asset_id is None else f"assets/{asset_id}"
     return f"{base_url}/{asset_path}"

src/entitysdk/types.py

@@ -32,6 +32,9 @@
 from entitysdk._server_schemas import EMCellMeshGenerationMethod as EMCellMeshGenerationMethod
 from entitysdk._server_schemas import EMCellMeshType as EMCellMeshType
 from entitysdk._server_schemas import EntityType as EntityType
+from entitysdk._server_schemas import MeasurableEntity as MeasurableEntity
+from entitysdk._server_schemas import MeasurementStatistic as MeasurementStatistic
+from entitysdk._server_schemas import MeasurementUnit as MeasurementUnit
 from entitysdk._server_schemas import ModifiedMorphologyMethodType as ModifiedMorphologyMethodType
 from entitysdk._server_schemas import PublicationType as PublicationType
 from entitysdk._server_schemas import Sex as Sex

tests/unit/conftest.py

@@ -1,4 +1,5 @@
 import uuid
+from typing import NamedTuple
 
 import pytest
 
@@ -7,6 +8,11 @@
 from tests.unit.util import PROJECT_ID, VIRTUAL_LAB_ID
 
 
+class Clients(NamedTuple):
+    with_context: Client
+    wout_context: Client
+
+
 @pytest.fixture(scope="session")
 def api_url():
     return "http://mock-host:8000"
@@ -46,6 +52,19 @@ def client(project_context, api_url, auth_token):
     return Client(api_url=api_url, project_context=project_context, token_manager=auth_token)
 
 
+@pytest.fixture
+def client_no_context(api_url, auth_token):
+    return Client(api_url=api_url, token_manager=auth_token)
+
+
+@pytest.fixture
+def clients(client, client_no_context):
+    return Clients(
+        with_context=client,
+        wout_context=client_no_context,
+    )
+
+
 @pytest.fixture
 def random_uuid():
     return uuid.uuid4()