Switch from pk to uuid for entity identifier (except users)

Author: edan-bainglass

aiida_restapi/repository/entity.py

@@ -17,8 +17,9 @@ class EntityRepository(t.Generic[EntityType, EntityModelType]):
     :ivar entity_cls: The AiiDA ORM entity class associated with this utility, e.g. `orm.User`, `orm.Node`, etc.
     """
 
-    def __init__(self, entity_class: type[EntityType]) -> None:
-        self.entity_class: type[EntityType] = entity_class
+    def __init__(self, entity_class: type[EntityType], identifier: t.Literal['pk', 'uuid']) -> None:
+        self.entity_class = entity_class
+        self.identifier = identifier
 
     def get_entity_schema(self, which: t.Literal['get', 'post'] | None = None) -> dict:
         """Get JSON schema for the AiiDA entity.
@@ -65,25 +66,25 @@ def get_entities(self, queries: QueryParams) -> PaginatedResults:
             results=[self.to_model(result) for result in results],
         )
 
-    def get_entity_by_id(self, entity_id: int) -> EntityModelType:
+    def get_entity_by_id(self, identifier: str | int) -> EntityModelType:
         """Get an AiiDA entity by id.
 
-        :param entity_id: The id of the entity to retrieve.
+        :param identifier: The id of the entity to retrieve.
         :return: The AiiDA entity model, e.g. `orm.User.Model`, `orm.Node.Model`, etc.
         """
-        entity = self.entity_class.collection.get(pk=entity_id)
+        entity = self.entity_class.collection.get(**{self.identifier: identifier})
         return self.to_model(entity)
 
-    def get_entity_extras(self, entity_id: int) -> dict[str, t.Any]:
+    def get_entity_extras(self, identifier: str | int) -> dict[str, t.Any]:
         """Get the extras of an entity.
 
-        :param entity_id: The id of the entity to retrieve the extras for.
+        :param identifier: The id of the entity to retrieve the extras for.
         :return: A dictionary with the entity extras.
         """
         return t.cast(
             dict,
             self.entity_class.collection.query(
-                filters={'pk': entity_id},
+                filters={self.identifier: identifier},
                 project=['extras'],
             ).first()[0],
         )

aiida_restapi/repository/node.py

@@ -74,13 +74,13 @@ def get_all_download_formats(self, full_type: str | None = None) -> dict:
 
         return all_formats
 
-    def get_node_repository_metadata(self, node_id: int) -> dict[str, dict]:
+    def get_node_repository_metadata(self, uuid: str) -> dict[str, dict]:
         """Get the repository metadata of a node.
 
-        :param node_id: The id of the node to retrieve the repository metadata for.
+        :param uuid: The uuid of the node to retrieve the repository metadata for.
         :return: A dictionary with the repository file metadata.
         """
-        node = self.entity_class.collection.get(pk=node_id)
+        node = self.entity_class.collection.get(uuid=uuid)
         total_size = 0
 
         def get_metadata(objects: list[File], path: str | None = None) -> dict[str, dict]:
@@ -113,7 +113,7 @@ def get_metadata(objects: list[File], path: str | None = None) -> dict[str, dict
                         'type': 'FILE',
                         'binary': binary,
                         'size': size,
-                        'download': f'/nodes/{node_id}/repo/contents?filename={obj_name}',
+                        'download': f'/nodes/{uuid}/repo/contents?filename={obj_name}',
                     }
                     total_size += size
 
@@ -126,39 +126,39 @@ def get_metadata(objects: list[File], path: str | None = None) -> dict[str, dict
                 'type': 'FILE',
                 'binary': True,
                 'size': total_size,
-                'download': f'/nodes/{node_id}/repo/contents',
+                'download': f'/nodes/{uuid}/repo/contents',
             }
 
         return metadata
 
-    def get_node_attributes(self, node_id: int) -> dict[str, t.Any]:
+    def get_node_attributes(self, uuid: str) -> dict[str, t.Any]:
         """Get the attributes of a node.
 
-        :param node_id: The id of the node to retrieve the attributes for.
+        :param uuid: The uuid of the node to retrieve the attributes for.
         :return: A dictionary with the node attributes.
         """
         return t.cast(
             dict,
             self.entity_class.collection.query(
-                filters={'pk': node_id},
+                filters={'uuid': uuid},
                 project=['attributes'],
             ).first()[0],
         )
 
     def get_node_links(
         self,
-        node_id: int,
+        uuid: str,
         queries: QueryParams,
         direction: t.Literal['incoming', 'outgoing'],
     ) -> PaginatedResults[NodeLinks]:
         """Get the incoming links of a node.
 
-        :param node_id: The id of the node to retrieve the incoming links for.
+        :param uuid: The uuid of the node to retrieve the incoming links for.
         :param queries: The query parameters, including filters, order_by, page_size, and page.
         :param direction: Specify whether to retrieve incoming or outgoing links.
         :return: The paginated requested linked nodes.
         """
-        node = self.entity_class.collection.get(pk=node_id)
+        node = self.entity_class.collection.get(uuid=uuid)
 
         start, end = (
             queries.page_size * (queries.page - 1),

aiida_restapi/routers/computers.py

@@ -18,7 +18,10 @@
 read_router = APIRouter()
 write_router = APIRouter()
 
-repository = EntityRepository[orm.Computer, orm.Computer.Model](orm.Computer)
+repository = EntityRepository[orm.Computer, orm.Computer.Model](
+    entity_class=orm.Computer,
+    identifier='pk',
+)
 
 
 @read_router.get('/computers/schema')
@@ -32,12 +35,15 @@ async def get_computers_schema(
 
     :param which: The type of schema to retrieve: 'get' or 'post'.
     :return: A dictionary with 'get' and 'post' keys containing the respective JSON schemas.
-    :raises HTTPException: 422 if the 'which' parameter is not 'get' or 'post'.
+    :raises HTTPException: 422 if the 'which' parameter is not 'get' or 'post',
+        500 for any other failures.
     """
     try:
         return repository.get_entity_schema(which=which)
-    except ValueError as err:
-        raise HTTPException(status_code=422, detail=str(err)) from err
+    except ValueError as exception:
+        raise HTTPException(status_code=422, detail=str(exception)) from exception
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception
 
 
 @read_router.get('/computers/projectable_properties', response_model=list[str])
@@ -68,39 +74,45 @@ async def get_computers(
 
 
 @read_router.get(
-    '/computers/{computer_id}',
+    '/computers/{pk}',
     response_model=orm.Computer.Model,
     response_model_exclude_none=True,
     response_model_exclude_unset=True,
 )
 @with_dbenv()
-async def get_computer(computer_id: int) -> orm.Computer.Model:
-    """Get AiiDA computer by id.
+async def get_computer(pk: str) -> orm.Computer.Model:
+    """Get AiiDA computer by pk.
 
-    :param computer_id: The id of the AiiDA computer.
+    :param pk: The pk of the AiiDA computer.
     :return: The computer model.
-    :raises HTTPException: 404 if the computer with the given id does not exist.
+    :raises HTTPException: 404 if the computer with the given pk does not exist,
+        500 for any other failures.
     """
     try:
-        return repository.get_entity_by_id(computer_id)
-    except NotExistent:
-        raise HTTPException(status_code=404, detail=f'Could not find a Computer with id {computer_id}')
+        return repository.get_entity_by_id(pk)
+    except NotExistent as exception:
+        raise HTTPException(status_code=404, detail=str(exception)) from exception
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception
 
 
-@read_router.get('/computers/{computer_id}/metadata', response_model=dict[str, t.Any])
+@read_router.get('/computers/{pk}/metadata', response_model=dict[str, t.Any])
 @with_dbenv()
-async def get_computer_metadata(computer_id: int) -> dict[str, t.Any]:
-    """Get metadata of an AiiDA computer by id.
+async def get_computer_metadata(pk: str) -> dict[str, t.Any]:
+    """Get metadata of an AiiDA computer by pk.
 
-    :param computer_id: The id of the AiiDA computer.
+    :param pk: The pk of the AiiDA computer.
     :return: The metadata dictionary of the computer.
-    :raises HTTPException: 404 if the computer with the given id does not exist.
+    :raises HTTPException: 404 if the computer with the given pk does not exist,
+        500 for any other failures.
     """
     try:
-        computer = repository.get_entity_by_id(computer_id)
+        computer = repository.get_entity_by_id(pk)
         return computer.metadata
-    except NotExistent:
-        raise HTTPException(status_code=404, detail=f'Could not find a Computer with id {computer_id}')
+    except NotExistent as exception:
+        raise HTTPException(status_code=404, detail=str(exception)) from exception
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception
 
 
 @write_router.post(
@@ -123,5 +135,5 @@ async def create_computer(
     """
     try:
         return repository.create_entity(computer_model)
-    except Exception as err:
-        raise HTTPException(status_code=500, detail=str(err))
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception))

aiida_restapi/routers/groups.py

@@ -18,7 +18,10 @@
 read_router = APIRouter()
 write_router = APIRouter()
 
-repository = EntityRepository[orm.Group, orm.Group.Model](orm.Group)
+repository = EntityRepository[orm.Group, orm.Group.Model](
+    entity_class=orm.Group,
+    identifier='uuid',
+)
 
 
 @read_router.get('/groups/schema')
@@ -32,12 +35,15 @@ async def get_groups_schema(
 
     :param which: The type of schema to retrieve: 'get' or 'post'.
     :return: A dictionary with 'get' and 'post' keys containing the respective JSON schemas.
-    :raises HTTPException: 422 if the 'which' parameter is not 'get' or 'post'.
+    :raises HTTPException: 422 if the 'which' parameter is not 'get' or 'post',
+        500 for any other failures.
     """
     try:
         return repository.get_entity_schema(which=which)
-    except ValueError as err:
-        raise HTTPException(status_code=422, detail=str(err)) from err
+    except ValueError as exception:
+        raise HTTPException(status_code=422, detail=str(exception)) from exception
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception
 
 
 @read_router.get('/groups/projectable_properties', response_model=list[str])
@@ -68,47 +74,47 @@ async def get_groups(
 
 
 @read_router.get(
-    '/groups/{group_id}',
+    '/groups/{uuid}',
     response_model=orm.Group.Model,
     response_model_exclude_none=True,
     response_model_exclude_unset=True,
 )
 @with_dbenv()
-async def get_group(group_id: int) -> orm.Group.Model:
-    """Get AiiDA group by id.
+async def get_group(uuid: str) -> orm.Group.Model:
+    """Get AiiDA group by uuid.
 
-    :param group_id: The id of the group to retrieve.
+    :param uuid: The uuid of the group to retrieve.
     :return: The AiiDA group model.
-    :raises HTTPException: 404 if a group with the given id does not exist,
+    :raises HTTPException: 404 if a group with the given uuid does not exist,
         500 for any other server error.
     """
     try:
-        return repository.get_entity_by_id(group_id)
-    except NotExistent:
-        raise HTTPException(status_code=404, detail=f'Could not find a Group with id {group_id}')
-    except Exception as err:
-        raise HTTPException(status_code=500, detail=str(err))
+        return repository.get_entity_by_id(uuid)
+    except NotExistent as exception:
+        raise HTTPException(status_code=404, detail=str(exception)) from exception
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception
 
 
 @read_router.get(
-    '/groups/{group_id}/extras',
+    '/groups/{uuid}/extras',
     response_model=dict[str, t.Any],
 )
 @with_dbenv()
-async def get_group_extras(group_id: int) -> dict[str, t.Any]:
+async def get_group_extras(uuid: str) -> dict[str, t.Any]:
     """Get the extras of a group.
 
-    :param group_id: The id of the group to retrieve the extras for.
+    :param uuid: The uuid of the group to retrieve the extras for.
     :return: A dictionary with the group extras.
-    :raises HTTPException: 404 if the group with the given id does not exist,
+    :raises HTTPException: 404 if the group with the given uuid does not exist,
         500 for other failures during retrieval.
     """
     try:
-        return repository.get_entity_extras(group_id)
-    except NotExistent:
-        raise HTTPException(status_code=404, detail=f'Could not find any group with id {group_id}')
-    except Exception as err:
-        raise HTTPException(status_code=500, detail=str(err)) from err
+        return repository.get_entity_extras(uuid)
+    except NotExistent as exception:
+        raise HTTPException(status_code=404, detail=str(exception)) from exception
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception
 
 
 @write_router.post(
@@ -131,5 +137,5 @@ async def create_group(
     """
     try:
         return repository.create_entity(group_model)
-    except Exception as err:
-        raise HTTPException(status_code=500, detail=str(err))
+    except Exception as exception:
+        raise HTTPException(status_code=500, detail=str(exception)) from exception