Leverage aiida-core pydantic models

Author: edan-bainglass

aiida_restapi/routers/comments.py aiida_restapi/common/__init__.pyaiida_restapi/routers/comments.py renamed to aiida_restapi/common/__init__.py

@@ -0,0 +1,19 @@
+"""Pagination utilities."""
+
+from __future__ import annotations
+
+import typing as t
+
+import pydantic as pdt
+from aiida.orm import Entity
+
+ResultType = t.TypeVar('ResultType', Entity.Model, t.Any)
+
+__all__ = ('PaginatedResults',)
+
+
+class PaginatedResults(pdt.BaseModel, t.Generic[ResultType]):
+    total: int
+    page: int
+    page_size: int
+    results: list[ResultType]

aiida_restapi/common/pagination.py

@@ -0,0 +1,90 @@
+"""REST API query utilities."""
+
+from __future__ import annotations
+
+import json
+import typing as t
+
+import pydantic as pdt
+from fastapi import HTTPException, Query
+
+
+class QueryParams(pdt.BaseModel):
+    filters: dict[str, t.Any] = pdt.Field(
+        default_factory=dict,
+        description='AiiDA QueryBuilder filters',
+        examples=[
+            {'node_type': {'==': 'data.core.int.Int.'}},
+            {'attributes.value': {'>': 42}},
+        ],
+    )
+    order_by: str | list[str] | dict[str, t.Any] | None = pdt.Field(
+        None,
+        description='Fields to sort by',
+        examples=[
+            {'attributes.value': 'desc'},
+        ],
+    )
+    page_size: pdt.PositiveInt = pdt.Field(
+        10,
+        description='Number of results per page',
+        examples=[10],
+    )
+    page: pdt.PositiveInt = pdt.Field(
+        1,
+        description='Page number',
+        examples=[1],
+    )
+
+
+def query_params(
+    filters: str | None = Query(
+        None,
+        description='AiiDA QueryBuilder filters as JSON string',
+    ),
+    order_by: str | None = Query(
+        None,
+        description='Comma-separated list of fields to sort by',
+    ),
+    page_size: pdt.PositiveInt = Query(
+        10,
+        description='Number of results per page',
+    ),
+    page: pdt.PositiveInt = Query(
+        1,
+        description='Page number',
+    ),
+) -> QueryParams:
+    """Parse query parameters into a structured object.
+
+    :param filters: AiiDA QueryBuilder filters as JSON string.
+    :param order_by: Comma-separated string of fields to sort by.
+    :param page_size: Number of results per page.
+    :param page: Page number.
+    :return: Structured query parameters.
+    :raises HTTPException: If filters cannot be parsed as JSON.
+    """
+    query_filters: dict[str, t.Any] = {}
+    query_order_by: str | list[str] | dict[str, t.Any] | None = None
+    if filters:
+        try:
+            query_filters = json.loads(filters)
+        except Exception as exception:
+            raise HTTPException(
+                status_code=400,
+                detail=f'Could not parse filters as JSON: {exception}',
+            ) from exception
+    if order_by:
+        try:
+            query_order_by = json.loads(order_by)
+        except Exception as exception:
+            raise HTTPException(
+                status_code=400,
+                detail=f'Could not parse order_by as JSON: {exception}',
+            ) from exception
+    return QueryParams(
+        filters=query_filters,
+        order_by=query_order_by,
+        page_size=page_size,
+        page=page,
+    )

aiida_restapi/common/query.py

@@ -0,0 +1,13 @@
+"""Common type variables."""
+
+from __future__ import annotations
+
+import typing as t
+
+from aiida import orm
+
+EntityType = t.TypeVar('EntityType', bound='orm.Entity')
+EntityModelType = t.TypeVar('EntityModelType', bound='orm.Entity.Model')
+
+NodeType = t.TypeVar('NodeType', bound='orm.Node')
+NodeModelType = t.TypeVar('NodeModelType', bound='orm.Node.Model')

aiida_restapi/common/types.py

@@ -6,7 +6,7 @@
 import graphene as gr
 from aiida.orm import Comment
 
-from aiida_restapi.filter_syntax import parse_filter_str
+from aiida_restapi.graphql.filter_syntax import parse_filter_str
 
 from .orm_factories import (
     ENTITY_DICT_TYPE,

aiida_restapi/aiida_db_mappings.py …ida_restapi/graphql/aiida_db_mappings.pyaiida_restapi/aiida_db_mappings.py renamed to aiida_restapi/graphql/aiida_db_mappings.py aiida_restapi/aiida_db_mappings.py renamed to aiida_restapi/graphql/aiida_db_mappings.py

@@ -6,7 +6,7 @@
 import graphene as gr
 from aiida.orm import Computer
 
-from aiida_restapi.filter_syntax import parse_filter_str
+from aiida_restapi.graphql.filter_syntax import parse_filter_str
 from aiida_restapi.graphql.plugins import QueryPlugin
 
 from .nodes import NodesQuery

aiida_restapi/graphql/comments.py

@@ -13,8 +13,8 @@
 
 from lark import Lark, Token, Tree
 
-from . import static
-from .utils import parse_date
+from .. import static
+from ..utils import parse_date
 
 FILTER_GRAMMAR = resources.open_text(static, 'filter_grammar.lark')
 

aiida_restapi/graphql/computers.py

@@ -6,7 +6,7 @@
 import graphene as gr
 from aiida.orm import Group
 
-from aiida_restapi.filter_syntax import parse_filter_str
+from aiida_restapi.graphql.filter_syntax import parse_filter_str
 from aiida_restapi.graphql.nodes import NodesQuery
 from aiida_restapi.graphql.plugins import QueryPlugin
 

aiida_restapi/filter_syntax.py aiida_restapi/graphql/filter_syntax.pyaiida_restapi/filter_syntax.py renamed to aiida_restapi/graphql/filter_syntax.py

@@ -6,7 +6,7 @@
 import graphene as gr
 from aiida.orm import Log
 
-from aiida_restapi.filter_syntax import parse_filter_str
+from aiida_restapi.graphql.filter_syntax import parse_filter_str
 
 from .orm_factories import (
     ENTITY_DICT_TYPE,