Modularize processing a DB response

A new data structure has been added to represent a gateway query
response: `GatewayQueryResponse`.
This new response changes the structure for the value of the top-level
`data` field by always returning a dictionary/JSON object. The keys of
the dictionary are the database `id`s and the value is a list of the
returned resources from each database.
In this way there is no need to alter/add anything to the returned
resources, which has been the case up to now, where the `id` for each
resource was prepended with the database `id`.
This is however still in place, somewhat, if a `QueryResource` is not
used for the query, but a straight-through result is desired. Then an
OPTIMADE-compliant response is produced, where the database `id` is
added as a `meta` entry to each returned resource.

To handle this new approach a new utility function for processing an
OPTIMADE database response has been written and implemented.

An Enum for OPTIMADE endpoint entry types has been implemented with
methods for returning the correct pydantic model for either the
resources or response (both many and single), removing the need to pass
the import path for these models.

Make entry ID ambiguity a warning instead of an exception.

.ci/test_queries.json

@@ -7,8 +7,7 @@
     "query_parameters": {
       "filter": "elements HAS \"Si\""
     },
-    "endpoint": "structures",
-    "endpoint_model": ["optimade.models.responses", "StructureResponseMany"]
+    "endpoint": "structures"
   },
   {
     "id": "twodbs_query_2",
@@ -18,8 +17,7 @@
     "query_parameters": {
       "filter": "elements HAS \"Si\""
     },
-    "endpoint": "structures",
-    "endpoint_model": ["optimade.models.responses", "StructureResponseMany"]
+    "endpoint": "structures"
   },
   {
     "id": "singledb_query_1",
@@ -30,7 +28,6 @@
       "filter": "elements HAS \"Si\"",
       "page_limit": 5
     },
-    "endpoint": "structures",
-    "endpoint_model": ["optimade.models.responses", "StructureResponseMany"]
+    "endpoint": "structures"
   }
 ]

optimade_gateway/common/utils.py

@@ -33,7 +33,7 @@ async def clean_python_types(data: Any) -> Any:
 
 
 def get_resource_attribute(
-    resource: Union[BaseModel, Dict[str, Any]],
+    resource: Union[BaseModel, Dict[str, Any], None],
     field: str,
     default: Any = None,
     disambiguate: bool = True,
@@ -53,7 +53,7 @@ def get_resource_attribute(
 
     Parameters:
         resource: The resource, from which to get the field value.
-        field: The resource field. This can be a comma-separated nested field, e.g.,
+        field: The resource field. This can be a dot-separated nested field, e.g.,
             `"attributes.base_url"`.
         default: The default value to return if `field` does not exist.
         disambiguate: Whether or not to "shortcut" a field value.
@@ -71,6 +71,9 @@ def get_resource_attribute(
         def _get_attr(mapping: dict, key: str, default: Any) -> Any:
             return mapping.get(key, default)
 
+    elif resource is None:
+        # Allow passing `None`, but simply return `default`
+        return default
     else:
         raise TypeError(
             "resource must be either a pydantic model or a Python dictionary, it was of type "

optimade_gateway/models/__init__.py

@@ -6,7 +6,7 @@
 
 from .databases import DatabaseCreate
 from .gateways import GatewayCreate, GatewayResource, GatewayResourceAttributes
-from .queries import QueryCreate, QueryResource, QueryState
+from .queries import GatewayQueryResponse, QueryCreate, QueryResource, QueryState
 from .resources import EntryResourceCreate
 from .responses import (
     DatabasesResponse,
@@ -25,6 +25,7 @@
     "DatabasesResponseSingle",
     "EntryResourceCreate",
     "GatewayCreate",
+    "GatewayQueryResponse",
     "GatewayResource",
     "GatewayResourceAttributes",
     "GatewaysResponse",

optimade_gateway/models/queries.py

@@ -1,13 +1,19 @@
 """Pydantic models/schemas for the Queries resource"""
 from enum import Enum
-from typing import Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 import warnings
 
 from optimade.models import (
     EntryResource,
     EntryResourceAttributes,
     EntryResponseMany,
     ErrorResponse,
+    ReferenceResource,
+    ReferenceResponseMany,
+    ReferenceResponseOne,
+    StructureResource,
+    StructureResponseMany,
+    StructureResponseOne,
 )
 from optimade.server.query_params import EntryListingQueryParams
 from pydantic import BaseModel, EmailStr, Field, validator
@@ -16,6 +22,39 @@
 from optimade_gateway.warnings import SortNotSupported
 
 
+class EndpointEntryType(Enum):
+    """Entry endpoint resource types, mapping to their pydantic models from the `optimade` package."""
+
+    REFERENCES = "references"
+    STRUCTURES = "structures"
+
+    def get_resource_model(self) -> Union[ReferenceResource, StructureResource]:
+        """Get the matching pydantic model for a resource."""
+        return {
+            "references": ReferenceResource,
+            "structures": StructureResource,
+        }[self.value]
+
+    def get_response_model(
+        self, single: bool = False
+    ) -> Union[
+        ReferenceResponseMany,
+        ReferenceResponseOne,
+        StructureResponseMany,
+        StructureResponseOne,
+    ]:
+        """Get the matching pydantic model for a successful response."""
+        if single:
+            return {
+                "references": ReferenceResponseOne,
+                "structures": StructureResponseOne,
+            }[self.value]
+        return {
+            "references": ReferenceResponseMany,
+            "structures": StructureResponseMany,
+        }[self.value]
+
+
 QUERY_PARAMETERS = EntryListingQueryParams()
 """Entry listing URL query parameters from the `optimade` package
 ([`EntryListingQueryParams`](https://www.optimade.org/optimade-python-tools/api_reference/server/query_params/#optimade.server.query_params.EntryListingQueryParams))."""
@@ -99,8 +138,17 @@ class QueryState(Enum):
     FINISHED = "finished"
 
 
+class GatewayQueryResponse(EntryResponseMany):
+    """Response from a Gateway Query."""
+
+    data: Dict[str, Union[List[EntryResource], List[Dict[str, Any]]]] = Field(
+        ...,
+        uniqueItems=True,
+    )
+
+
 class QueryResourceAttributes(EntryResourceAttributes):
-    """Attributes for an OPTIMADE gateway query"""
+    """Attributes for an OPTIMADE gateway query."""
 
     gateway_id: str = Field(
         ...,
@@ -117,45 +165,26 @@ class QueryResourceAttributes(EntryResourceAttributes):
         title="State",
         type="enum",
     )
-    response: Optional[Union[EntryResponseMany, ErrorResponse]] = Field(
+    response: Optional[Union[GatewayQueryResponse, ErrorResponse]] = Field(
         None,
         description="Response from gateway query.",
         type="object",
     )
-    endpoint: str = Field(
-        ..., description="The entry endpoint queried, e.g., 'structures'."
-    )
-    endpoint_model: Tuple[str, str] = Field(
-        ...,
-        description=(
-            "The full importable path to the pydantic response model class (not an instance of "
-            "the class). It should be a tuple of the Python module and the Class name."
-        ),
+    endpoint: EndpointEntryType = Field(
+        EndpointEntryType.STRUCTURES,
+        description="The entry endpoint queried, e.g., 'structures'.",
+        title="Endpoint",
+        type="enum",
     )
 
     @validator("endpoint")
-    def remove_endpoints_slashes(cls, value: str) -> str:
-        """Remove prep-/appended slashes (`/`)"""
-        org_value = value
-        value = value.strip()
-        while value.startswith("/"):
-            value = value[1:]
-        while value.endswith("/"):
-            value = value[:-1]
-        if not value:
-            raise ValueError(
-                "endpoint must not be an empty string or be prep-/appended with slashes (`/`). "
-                f"Original value: {org_value!r}. Final value (after removing prep-/appended "
-                f"slashes): {value!r}"
-            )
-
-        # Temporarily only allow queries to "structures" endpoints.
-        if value != "structures":
+    def only_allow_structures(cls, value: EndpointEntryType) -> EndpointEntryType:
+        """Temporarily only allow queries to "structures" endpoints."""
+        if value != EndpointEntryType.STRUCTURES:
             raise NotImplementedError(
                 'OPTIMADE Gateway temporarily only supports queries to "structures" endpoints, '
                 'i.e.: endpoint="structures"'
             )
-
         return value
 
 
@@ -175,8 +204,7 @@ class QueryCreate(EntryResourceCreate, QueryResourceAttributes):
     """Model for creating new Query resources in the MongoDB"""
 
     state: Optional[QueryState]
-    endpoint: Optional[str]
-    endpoint_model: Optional[Tuple[str, str]]
+    endpoint: Optional[EndpointEntryType]
 
     @validator("query_parameters")
     def sort_not_supported(