Add relationships functionality

optimade/models/entries.py

@@ -1,7 +1,7 @@
 # pylint: disable=line-too-long
 from datetime import datetime
-from typing import Optional, Dict, List
-from pydantic import BaseModel, Schema
+from typing import Optional, Dict, List, Union
+from pydantic import BaseModel, Schema, validator
 
 from .jsonapi import Relationships, Attributes, Resource, Relationship
 
@@ -15,11 +15,33 @@
 )
 
 
+class TypedRelationship(Relationship):
+    @validator("data", whole=True)
+    def check_rel_type(cls, data, values):
+        if hasattr(cls, "_req_type") and any(obj.type != cls._req_type for obj in data):
+            raise ValueError("Object stored in relationship data has wrong type")
+        return data
+
+
+class ReferenceRelationship(TypedRelationship):
+    _req_type = "references"
+
+
+class StructureRelationship(TypedRelationship):
+    _req_type = "structures"
+
+
 class EntryRelationships(Relationships):
     """This model wraps the JSON API Relationships to include type-specific top level keys. """
 
-    references: Optional[Relationship] = Schema(
-        ..., description="Object containing links to relationships with other entries."
+    references: Optional[ReferenceRelationship] = Schema(
+        ...,
+        description="Object containing links to relationships with entries of the `references` type.",
+    )
+
+    structures: Optional[StructureRelationship] = Schema(
+        ...,
+        description="Object containing links to relationships with entries of the `structures` type.",
     )
 
 

optimade/models/jsonapi.py

@@ -139,7 +139,7 @@ class RelationshipLinks(BaseModel):
     @validator("related", always=True)
     def either_self_or_related_must_be_specified(cls, v, values):
         if values.get("self", None) is None and v is None:
-            raise AssertionError(
+            raise ValueError(
                 "Either 'self' or 'related' MUST be specified for RelationshipLinks"
             )
         return v
@@ -167,7 +167,7 @@ def at_least_one_relationship_key_must_be_set(cls, v, values):
             and values.get("data", None) is None
             and v is None
         ):
-            raise AssertionError(
+            raise ValueError(
                 "Either 'links', 'data', or 'meta' MUST be specified for relationship"
             )
         return v
@@ -186,7 +186,7 @@ class Relationships(BaseModel):
 
     @validator("id", "type")
     def check_illegal_relationships_fields(cls, v):
-        raise AssertionError('"id", "type" MUST NOT be fields under relationships')
+        raise ValueError('"id", "type" MUST NOT be fields under relationships')
 
 
 class ResourceLinks(BaseModel):
@@ -271,7 +271,7 @@ def either_data_meta_or_errors_must_be_set(cls, v, values):
             and values.get("meta", None) is None
             and v is None
         ):
-            raise AssertionError(
+            raise ValueError(
                 "Either 'data', 'meta', or 'errors' must be specified in the top-level response"
             )
         return v

optimade/models/tests/test_models.py

@@ -50,6 +50,12 @@ def test_simple_relationships(self):
         }
         EntryRelationships(**relationship)
 
+        relationship = {
+            "references": {"data": [{"id": "Dijkstra1968", "type": "structures"}]}
+        }
+        with self.assertRaises(ValidationError):
+            EntryRelationships(**relationship)
+
 
 def test_constrained_list():
     class ConListModel(BaseModel):

optimade/server/entry_collections.py

@@ -119,19 +119,18 @@ def find(
             fields = set(params.response_fields.split(","))
         else:
             fields = all_fields.copy()
-        results = []
 
-        # build up set of included data with unique IDs
-        included = []
-        id_set = set()
+        results = []
+        included = {}
         for doc in self.collection.find(**criteria):
             results.append(self.resource_cls(**self.resource_mapper.map_back(doc)))
-            for reference in (
-                doc.get("relationships", {}).get("references", {}).get("data", [])
-            ):
-                if reference["id"] not in id_set:
-                    included.append(reference)
-                    id_set.add(reference["id"])
+            # collapse list of references into dict by ID and take only one per ID
+            refs = doc.get("relationships", {}).get("references", {}).get("data", [])
+            for ref in refs:
+                # could check here and raise a warning if any IDs clash
+                included[ref["id"]] = ref
+
+        included = list(included.values())
 
         if isinstance(params, SingleEntryQueryParams):
             results = results[0] if results else None