Added logic and tests for HAS ONLY filtering on relationships

Materials-Consortia · Mar 30, 2020 · 03dee04 · 03dee04
1 parent 0e49c16
commit 03dee04
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 3 deletions.
diff --git a/optimade/filtertransformers/mongo.py b/optimade/filtertransformers/mongo.py
@@ -44,7 +44,9 @@ def postprocess(self, query):
         return query
 
     def transform(self, tree):
-        return self.postprocess(super().transform(tree))
+        res = self.postprocess(super().transform(tree))
+        print(res)
+        return res
 
     def filter(self, arg):
         # filter: expression*
@@ -362,8 +364,23 @@ def check_for_entry_type(prop, expr):
 
         def replace_with_relationship(subdict, prop, expr):
             _prop, _field = str(prop).split(".")
-            subdict[f"relationships.{_prop}.data.{_field}"] = expr
+
+            # in the case of HAS ONLY, the size operator needs to be applied
+            # one level up, i.e. excluding the field
+            if "$size" in expr:
+                if "$and" not in subdict:
+                    subdict["$and"] = []
+                subdict["$and"].extend(
+                    [
+                        {f"relationships.{_prop}.data": {"$size": expr.pop("$size")}},
+                        {f"relationships.{_prop}.data.{_field}": expr},
+                    ]
+                )
+            else:
+                subdict[f"relationships.{_prop}.data.{_field}"] = expr
+
             subdict.pop(prop)
+
             return subdict
 
         return recursive_postprocessing(

diff --git a/tests/filtertransformers/test_mongo.py b/tests/filtertransformers/test_mongo.py
@@ -321,7 +321,33 @@ def test_filtering_on_relationships(self):
 
         self.assertEqual(
             self.transform('structures.id HAS ONLY "dummy/2019"'),
-            {"relationships.structures.data.id": {"$all": ["dummy/2019"], "$size": 1}},
+            {
+                "$and": [
+                    {"relationships.structures.data": {"$size": 1}},
+                    {"relationships.structures.data.id": {"$all": ["dummy/2019"]}},
+                ]
+            },
+        )
+
+        self.assertEqual(
+            self.transform(
+                'structures.id HAS ONLY "dummy/2019" AND structures.id HAS "dummy/2019"'
+            ),
+            {
+                "$and": [
+                    {
+                        "$and": [
+                            {"relationships.structures.data": {"$size": 1}},
+                            {
+                                "relationships.structures.data.id": {
+                                    "$all": ["dummy/2019"]
+                                }
+                            },
+                        ]
+                    },
+                    {"relationships.structures.data.id": {"$in": ["dummy/2019"]}},
+                ],
+            },
         )
 
     def test_not_implemented(self):