Feature/computed field adjustments (#107)

# PR Context
- uses robert-koch-institut/mex-common#230

# Added
- add support for computed fields in graph queries

# Changes

- BREAKING: make `MEX_EXTRACTED_PRIMARY_SOURCE` an instance of its own
class
instead of ExtractedPrimarySource in order to set static provenance
identifiers

CHANGELOG.md

@@ -9,8 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- add support for computed fields in graph queries
+
 ### Changes
 
+- BREAKING: make `MEX_EXTRACTED_PRIMARY_SOURCE` an instance of its own class
+  instead of ExtractedPrimarySource in order to set static provenance identifiers
+
 ### Deprecated
 
 ### Removed

mex/backend/fields.py

@@ -1,22 +1,15 @@
 from collections.abc import Callable, Generator, Mapping
 from types import NoneType, UnionType
-from typing import (
-    Annotated,
-    Any,
-    Union,
-    get_args,
-    get_origin,
-)
-
-from pydantic import BaseModel
-from pydantic.fields import FieldInfo
+from typing import Annotated, Any, Union, get_args, get_origin
 
 from mex.common.models import (
     ADDITIVE_MODEL_CLASSES_BY_NAME,
     EXTRACTED_MODEL_CLASSES_BY_NAME,
     PREVENTIVE_MODEL_CLASSES_BY_NAME,
     SUBTRACTIVE_MODEL_CLASSES_BY_NAME,
+    BaseModel,
 )
+from mex.common.models.base import GenericFieldInfo
 from mex.common.types import MERGED_IDENTIFIER_CLASSES, Link, LiteralStringType, Text
 
 
@@ -38,14 +31,14 @@ def _get_inner_types(annotation: Any) -> Generator[type, None, None]:
         yield annotation
 
 
-def _contains_only_types(field: FieldInfo, *types: type) -> bool:
+def _contains_only_types(field: GenericFieldInfo, *types: type) -> bool:
     """Return whether a `field` is annotated as one of the given `types`.
 
     Unions, lists and type annotations are checked for their inner types and only the
     non-`NoneType` types are considered for the type-check.
 
     Args:
-        field: A pydantic `FieldInfo` object
+        field: A `GenericFieldInfo` instance
         types: Types to look for in the field's annotation
 
     Returns:
@@ -58,7 +51,7 @@ def _contains_only_types(field: FieldInfo, *types: type) -> bool:
 
 def _group_fields_by_class_name(
     model_classes_by_name: Mapping[str, type[BaseModel]],
-    predicate: Callable[[FieldInfo], bool],
+    predicate: Callable[[GenericFieldInfo], bool],
 ) -> dict[str, list[str]]:
     """Group the field names by model class and filter them by the given predicate.
 
@@ -73,7 +66,7 @@ def _group_fields_by_class_name(
         name: sorted(
             {
                 field_name
-                for field_name, field_info in cls.model_fields.items()
+                for field_name, field_info in cls.get_all_fields().items()
                 if predicate(field_info)
             }
         )
@@ -144,7 +137,7 @@ def _group_fields_by_class_name(
     name: sorted(
         {
             field_name
-            for field_name in cls.model_fields
+            for field_name in cls.get_all_fields()
             if field_name
             not in (
                 *FROZEN_FIELDS_BY_CLASS_NAME[name],
@@ -162,7 +155,7 @@ def _group_fields_by_class_name(
     name: sorted(
         {
             field_name
-            for field_name in cls.model_fields
+            for field_name in cls.get_all_fields()
             if field_name in FROZEN_FIELDS_BY_CLASS_NAME[name]
             and field_name
             not in (

mex/backend/graph/connector.py

@@ -1,8 +1,9 @@
 import json
 from string import Template
-from typing import Any
+from typing import Annotated, Any, Literal, cast
 
 from neo4j import Driver, GraphDatabase
+from pydantic import Field
 
 from mex.backend.fields import (
     FINAL_FIELDS_BY_CLASS_NAME,
@@ -15,9 +16,7 @@
 )
 from mex.backend.graph.models import Result
 from mex.backend.graph.query import QueryBuilder
-from mex.backend.graph.transform import (
-    expand_references_in_search_result,
-)
+from mex.backend.graph.transform import expand_references_in_search_result
 from mex.backend.settings import BackendSettings
 from mex.backend.transform import to_primitive
 from mex.common.connector import BaseConnector
@@ -33,17 +32,35 @@
     AnyRuleModel,
     ExtractedPrimarySource,
 )
+from mex.common.models.primary_source import BasePrimarySource
 from mex.common.transform import ensure_prefix, to_key_and_values
-from mex.common.types import AnyPrimitiveType, Identifier, Link, Text
-
-MEX_EXTRACTED_PRIMARY_SOURCE = ExtractedPrimarySource.model_construct(
-    hadPrimarySource=MEX_PRIMARY_SOURCE_STABLE_TARGET_ID,
-    identifier=MEX_PRIMARY_SOURCE_IDENTIFIER,
-    identifierInPrimarySource=MEX_PRIMARY_SOURCE_IDENTIFIER_IN_PRIMARY_SOURCE,
-    stableTargetId=MEX_PRIMARY_SOURCE_STABLE_TARGET_ID,
+from mex.common.types import (
+    AnyPrimitiveType,
+    ExtractedPrimarySourceIdentifier,
+    Identifier,
+    Link,
+    MergedPrimarySourceIdentifier,
+    Text,
 )
 
 
+class MExPrimarySource(BasePrimarySource):
+    """An automatically extracted metadata set describing a primary source."""
+
+    entityType: Annotated[
+        Literal["ExtractedPrimarySource"], Field(alias="$type", frozen=True)
+    ] = "ExtractedPrimarySource"
+    hadPrimarySource: MergedPrimarySourceIdentifier = (
+        MEX_PRIMARY_SOURCE_STABLE_TARGET_ID
+    )
+    identifier: ExtractedPrimarySourceIdentifier = MEX_PRIMARY_SOURCE_IDENTIFIER
+    identifierInPrimarySource: str = MEX_PRIMARY_SOURCE_IDENTIFIER_IN_PRIMARY_SOURCE
+    stableTargetId: MergedPrimarySourceIdentifier = MEX_PRIMARY_SOURCE_STABLE_TARGET_ID
+
+
+MEX_EXTRACTED_PRIMARY_SOURCE = MExPrimarySource()
+
+
 class GraphConnector(BaseConnector):
     """Connector to handle authentication and transactions with the graph database."""
 
@@ -112,7 +129,7 @@ def _seed_indices(self) -> Result:
 
     def _seed_data(self) -> list[Identifier]:
         """Ensure the primary source `mex` is seeded and linked to itself."""
-        return self.ingest([MEX_EXTRACTED_PRIMARY_SOURCE])
+        return self.ingest([cast(ExtractedPrimarySource, MEX_EXTRACTED_PRIMARY_SOURCE)])
 
     def close(self) -> None:
         """Close the connector's underlying requests session."""

mex/backend/main.py

@@ -1,5 +1,6 @@
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
+from itertools import chain
 from typing import Any
 
 import uvicorn
@@ -19,7 +20,7 @@
 from mex.backend.settings import BackendSettings
 from mex.common.cli import entrypoint
 from mex.common.connector import CONNECTOR_STORE
-from mex.common.types import Identifier
+from mex.common.types import EXTRACTED_IDENTIFIER_CLASSES, MERGED_IDENTIFIER_CLASSES
 from mex.common.types.identifier import MEX_ID_PATTERN
 
 
@@ -45,12 +46,12 @@ def create_openapi_schema() -> dict[str, Any]:
         routes=app.routes,
         servers=[dict(url=settings.backend_api_url)],
     )
-    for subclass in Identifier.__subclasses__():
-        name = subclass.__name__
+    for identifier in chain(EXTRACTED_IDENTIFIER_CLASSES, MERGED_IDENTIFIER_CLASSES):
+        name = identifier.__name__
         openapi_schema["components"]["schemas"][name] = {
             "title": name,
             "type": "string",
-            "description": subclass.__doc__,
+            "description": identifier.__doc__,
             "pattern": MEX_ID_PATTERN,
         }