Improve handling of MongoDB ObjectID

optimade/filtertransformers/mongo.py

@@ -1,6 +1,7 @@
 import copy
 from lark import Transformer, v_args, Token
 from optimade.server.mappers import BaseResourceMapper
+from optimade.server.exceptions import BadRequest
 
 
 class MongoTransformer(Transformer):
@@ -41,6 +42,7 @@ def postprocess(self, query):
         query = self._apply_relationship_filtering(query)
         query = self._apply_length_operators(query)
         query = self._apply_unknown_or_null_filter(query)
+        query = self._apply_mongo_id_filter(query)
 
         return query
 
@@ -445,6 +447,35 @@ def replace_known_filter_with_or(subdict, prop, expr):
             filter_, check_for_known_filter, replace_known_filter_with_or
         )
 
+    def _apply_mongo_id_filter(self, filter_: dict) -> dict:
+        """This method loops through the query and replaces any operations
+        on the special Mongodb `_id` key with the corresponding operation
+        on a BSON `ObjectId` type.
+        """
+
+        def check_for_id_key(prop, _):
+            """ Find cases where the query dict is operating on the `_id` field. """
+            return prop == "_id"
+
+        def replace_str_id_with_objectid(subdict, prop, expr):
+            from bson import ObjectId
+
+            for operator in subdict[prop]:
+                val = subdict[prop][operator]
+                if operator not in ("$eq", "$ne"):
+                    if self.mapper is not None:
+                        prop = self.mapper.alias_of(prop)
+                    raise BadRequest(
+                        detail=f"Operator not supported for query on field {prop}, can only test for equality"
+                    )
+                if isinstance(val, str):
+                    subdict[prop][operator] = ObjectId(val)
+            return subdict
+
+        return recursive_postprocessing(
+            filter_, check_for_id_key, replace_str_id_with_objectid
+        )
+
 
 def recursive_postprocessing(filter_, condition, replacement):
     """Recursively descend into the query, checking each dictionary

optimade/server/entry_collections/entry_collections.py

@@ -171,9 +171,12 @@ def handle_query_params(
             cursor_kwargs["limit"] = CONFIG.page_limit
 
         cursor_kwargs["fields"] = self.all_fields
-        cursor_kwargs["projection"] = [
-            self.resource_mapper.alias_for(f) for f in self.all_fields
-        ]
+        cursor_kwargs["projection"] = {
+            f"{self.resource_mapper.alias_for(f)}": True for f in self.all_fields
+        }
+
+        if "_id" not in cursor_kwargs["projection"]:
+            cursor_kwargs["projection"]["_id"] = False
 
         if getattr(params, "sort", False):
             cursor_kwargs["sort"] = self.parse_sort_params(params.sort)

optimade/server/entry_collections/mongo.py

@@ -116,6 +116,8 @@ def find(
 
         results = []
         for doc in self.collection.find(**criteria):
+            if criteria.get("projection", {}).get("_id"):
+                doc["_id"] = str(doc["_id"])
             results.append(self.resource_cls(**self.resource_mapper.map_back(doc)))
 
         nresults_now = len(results)

optimade/server/mappers/entries.py

@@ -146,9 +146,6 @@ def map_back(cls, doc: dict) -> dict:
             A resource object in OPTIMADE format.
 
         """
-        if "_id" in doc:
-            del doc["_id"]
-
         mapping = ((real, alias) for alias, real in cls.all_aliases())
         newdoc = {}
         reals = {real for alias, real in cls.all_aliases()}

optimade_config.json

@@ -23,6 +23,7 @@
     "aliases": {
         "structures": {
             "id": "task_id",
+            "immutable_id": "_id",
             "chemical_formula_descriptive": "pretty_formula",
             "chemical_formula_reduced": "pretty_formula",
             "chemical_formula_anonymous": "formula_anonymous"

tests/filtertransformers/test_mongo.py

@@ -3,6 +3,7 @@
 from lark.exceptions import VisitError
 
 from optimade.filterparser import LarkParser, ParserError
+from optimade.server.exceptions import BadRequest
 
 
 class TestMongoTransformer:
@@ -378,6 +379,29 @@ def test_unaliased_length_operator(self):
             "cartesian_site_positions.11": {"$exists": True}
         }
 
+    def test_mongo_special_id(self, mapper):
+
+        from optimade.filtertransformers.mongo import MongoTransformer
+        from bson import ObjectId
+
+        class MyMapper(mapper("StructureMapper")):
+            ALIASES = (("immutable_id", "_id"),)
+
+        transformer = MongoTransformer(mapper=MyMapper())
+        parser = LarkParser(version=self.version, variant=self.variant)
+
+        assert transformer.transform(
+            parser.parse('immutable_id = "5cfb441f053b174410700d02"')
+        ) == {"_id": {"$eq": ObjectId("5cfb441f053b174410700d02")}}
+
+        assert transformer.transform(
+            parser.parse('immutable_id != "5cfb441f053b174410700d02"')
+        ) == {"_id": {"$ne": ObjectId("5cfb441f053b174410700d02")}}
+
+        for op in ("CONTAINS", "STARTS WITH", "ENDS WITH", "HAS"):
+            with pytest.raises(BadRequest):
+                transformer.transform(parser.parse(f'immutable_id {op} "abcdef"'))
+
     def test_aliased_length_operator(self, mapper):
         from optimade.filtertransformers.mongo import MongoTransformer