0.15.0

atlasq/queryset/index.py

@@ -1,7 +1,7 @@
 import fnmatch
 from enum import Enum
 from logging import getLogger
-from typing import Dict, List
+from typing import Dict, List, Union
 
 import requests
 from atlasq.queryset.exceptions import AtlasIndexError, AtlasIndexFieldError
@@ -123,22 +123,27 @@ def ensure_index_exists(
             self.ensured = False
         return self.ensured
 
-    def _set_indexed_fields(self, index_result: Dict, base_field: str = ""):
-        lucene_type = index_result["type"]
-        if lucene_type in [
-            AtlasIndexType.DOCUMENT.value,
-            AtlasIndexType.EMBEDDED_DOCUMENT.value,
-        ]:
-            if not index_result.get("dynamic", False):
-                for field, value in index_result.get("fields", {}).items():
-                    field = f"{base_field}.{field}" if base_field else field
-                    self._set_indexed_fields(value, base_field=field)
-            else:
-                self._indexed_fields[f"{base_field}.*" if base_field else "*"] = ""
-        if base_field:
-            if lucene_type not in AtlasIndexType.values():
-                logger.warning(f"Lucene type {lucene_type} not configured")
-            self._indexed_fields[base_field] = lucene_type
+    def _set_indexed_fields(self, index_result: Union[Dict, List], base_field: str = ""):
+        if isinstance(index_result, list):
+            for obj in index_result:
+                self._set_indexed_fields(obj, base_field=base_field)
+        else:
+            lucene_type = index_result["type"]
+            if lucene_type in [
+                AtlasIndexType.DOCUMENT.value,
+                AtlasIndexType.EMBEDDED_DOCUMENT.value,
+            ]:
+                if not index_result.get("dynamic", False):
+                    for field, value in index_result.get("fields", {}).items():
+                        field = f"{base_field}.{field}" if base_field else field
+                        self._set_indexed_fields(value, base_field=field)
+                else:
+                    self._indexed_fields[f"{base_field}.*" if base_field else "*"] = ""
+            if base_field:
+                if lucene_type not in AtlasIndexType.values():
+                    logger.warning(f"Lucene type {lucene_type} not configured")
+                else:
+                    self._indexed_fields[base_field] = lucene_type
 
     def _set_indexed_from_mappings(self, index_result: Dict):
         mappings = index_result["mappings"]

atlasq/queryset/transform.py

@@ -1,5 +1,6 @@
 import datetime
 import logging
+import re
 from typing import Any, Dict, List, Tuple, Union
 
 from atlasq.queryset.exceptions import AtlasFieldError, AtlasIndexFieldError
@@ -10,6 +11,22 @@
 logger = logging.getLogger(__name__)
 
 
+def mergedicts(dict1, dict2):
+    for k in set(dict1.keys()).union(dict2.keys()):
+        if k in dict1 and k in dict2:
+            if isinstance(dict1[k], dict) and isinstance(dict2[k], dict):
+                yield k, dict(mergedicts(dict1[k], dict2[k]))
+            else:
+                # If one of the values is not a dict, you can't continue merging it.
+                # Value from second dict overrides one in first and we move on.
+                yield k, dict2[k]
+                # Alternatively, replace this with exception raiser to alert you of value conflicts
+        elif k in dict1:
+            yield k, dict1[k]
+        else:
+            yield k, dict2[k]
+
+
 class AtlasTransform:
 
     id_keywords = [
@@ -35,6 +52,8 @@ class AtlasTransform:
         "icontains",
         "startswith",
         "istartswith",
+        "iendswith",
+        "endswith",
         "iwholeword",
         "wholeword",
         "not",
@@ -51,23 +70,13 @@ class AtlasTransform:
     range_keywords = ["gt", "gte", "lt", "lte"]
     equals_keywords = []
     equals_type_supported = (bool, ObjectId, int, datetime.datetime)
-    text_keywords = [
-        "contains",
-        "icontains",
-        "iwholeword",
-        "wholeword",
-        "exact",
-        "iexact",
-        "eq",
-    ]
+    startswith_keywords = ["startswith", "istartswith"]
+    endswith_keywords = ["endswith", "iendswith"]
+    text_keywords = ["iwholeword", "wholeword", "exact", "iexact", "eq", "contains", "icontains"]
     all_keywords = ["all"]
     regex_keywords = ["regex", "iregex"]
     size_keywords = ["size"]
     not_converted = [
-        "istartswith",
-        "startswith",
-        "contains",
-        "icontains",
         "mod",
         "match",
     ]
@@ -147,9 +156,13 @@ def _single_equals(self, path: str, value: Union[ObjectId, bool]):
             }
         }
 
+    def _contains(self, path: str, value: Any, keyword: str = None):
+        if not keyword:
+            return {path: {"$elemMatch": value}}
+        return {path: {"$elemMatch": {f"${keyword}": value}}}
+
     def _equals(self, path: str, value: Union[List[Union[ObjectId, bool]], ObjectId, bool]) -> Dict:
         if isinstance(value, list):
-
             values = value
             if not values:
                 raise AtlasFieldError(f"Text search for equals on {path=} cannot be empty")
@@ -166,6 +179,16 @@ def _text(self, path: str, value: Any) -> Dict:
             "text": {"query": value, "path": path},
         }
 
+    def _startswith(self, path: str, value: Any) -> Dict:
+        if not value:
+            raise AtlasFieldError(f"Text search for {path} cannot be {value}")
+        return self._regex(path, f"{re.escape(value)}.*")
+
+    def _endswith(self, path: str, value: Any) -> Dict:
+        if not value:
+            raise AtlasFieldError(f"Text search for {path} cannot be {value}")
+        return self._regex(path, f".*{re.escape(value)}")
+
     def _size(self, path: str, value: int, operator: str) -> Dict:
         if not isinstance(value, int):
             raise AtlasFieldError(f"Size search for {path} must be an int")
@@ -230,17 +253,16 @@ def transform(self) -> Tuple[List[Dict], List[Dict], List[Dict]]:
         negative = []
 
         for key, value in self.atlas_query.items():
-            # if to_go is positive, we add the element in the positive list
-            # if to_go is negative, we add the element in the negative list
-            to_go = 1
+            # if the value is positive, we add the element in the positive list
+            # if the value is negative, we add the element in the negative list
+            positive = 1
             if isinstance(value, QuerySet):
                 logger.debug("Casting queryset to list, otherwise the aggregation will fail")
                 value = list(value)
             key_parts = key.split("__")
             obj = None
             path = ""
             for i, keyword in enumerate(key_parts):
-
                 if keyword in self.id_keywords:
                     keyword = "_id"
                     key_parts[i] = keyword
@@ -255,17 +277,17 @@ def transform(self) -> Tuple[List[Dict], List[Dict], List[Dict]]:
                 if keyword in self.not_converted:
                     raise NotImplementedError(f"Keyword {keyword} not implemented yet")
                 if keyword in self.negative_keywords:
-                    to_go *= -1
+                    positive *= -1
 
                 if keyword in self.size_keywords:
                     # it must the last keyword, otherwise we do not support it
                     if i != len(key_parts) - 1:
                         raise NotImplementedError(f"Keyword {keyword} not implemented yet")
-                    other_aggregations.append(self._size(path, value, "eq" if to_go == 1 else "ne"))
+                    other_aggregations.append(self._size(path, value, "eq" if positive == 1 else "ne"))
                     break
                 if keyword in self.exists_keywords:
                     if value is False:
-                        to_go *= -1
+                        positive *= -1
                     obj = self._exists(path)
                     break
 
@@ -284,8 +306,14 @@ def transform(self) -> Tuple[List[Dict], List[Dict], List[Dict]]:
                 if keyword in self.all_keywords:
                     obj = self._all(path, value)
                     break
+                if keyword in self.startswith_keywords:
+                    obj = self._startswith(path, value)
+                    break
+                if keyword in self.endswith_keywords:
+                    obj = self._endswith(path, value)
+                    break
                 if keyword in self.type_keywords:
-                    if to_go == -1:
+                    if positive == -1:
                         raise NotImplementedError(f"At the moment you can't have a negative `{keyword}` keyword")
                     other_aggregations.append(self._type(path, value))
             else:
@@ -297,13 +325,13 @@ def transform(self) -> Tuple[List[Dict], List[Dict], List[Dict]]:
                 if self.atlas_index.ensured:
                     self._ensure_path_is_indexed(path.split("."))
                 # we are wrapping the result to an embedded document
-                converted = self._convert_to_embedded_document(path.split("."), obj, positive=to_go == 1)
+                converted = self._convert_to_embedded_document(path.split("."), obj, positive=positive == 1)
                 if obj != converted:
                     # we have an embedded object
                     # the mustNot is done inside the embedded document clause
                     affirmative = self.merge_embedded_documents(converted, affirmative)
                 else:
-                    if to_go == 1:
+                    if positive == 1:
                         affirmative.append(converted)
                     else:
                         negative.append(converted)

tests/queryset/test_index.py

@@ -102,6 +102,9 @@ def test_set_indexed_fields(self):
             }
         )
         self.assertCountEqual(index._indexed_fields, ["*"])
+        index._indexed_fields.clear()
+        index._set_indexed_fields([{"type": "string"}, {"type": "number"}], "f")
+        self.assertCountEqual(index._indexed_fields, ["f"])
 
     def test_ensure_index_exists(self):
         index = AtlasIndex("myindex")

tests/queryset/test_transform.py

@@ -752,7 +752,7 @@ def test__equals_list_bool(self):
         self.assertEqual(res["compound"]["should"][0]["equals"]["path"], "field")
         self.assertEqual(res["compound"]["should"][0]["equals"]["value"], True)
 
-    def test(self):
+    def test_equal(self):
         q = AtlasQ(f=3)
         t = AtlasTransform(q.query, AtlasIndex("test"))
         res = t._text("field", "aaa")
@@ -768,6 +768,26 @@ def test_none(self):
         with self.assertRaises(AtlasFieldError):
             t._text("field", None)
 
+    def test_convert_startswith(self):
+        q = AtlasQ(f__startswith="test?")
+        t = AtlasTransform(q.query, AtlasIndex("test"))
+        res = t._startswith("f", "test?")
+        self.assertIn("regex", res)
+        self.assertIn("query", res["regex"])
+        self.assertIn("test\\?.*", res["regex"]["query"])
+        self.assertIn("path", res["regex"])
+        self.assertIn("f", res["regex"]["path"])
+
+    def test_convert_endswith(self):
+        q = AtlasQ(f__endswith="test?")
+        t = AtlasTransform(q.query, AtlasIndex("test"))
+        res = t._endswith("f", "test?")
+        self.assertIn("regex", res)
+        self.assertIn("query", res["regex"])
+        self.assertIn(".*test\\?", res["regex"]["query"])
+        self.assertIn("path", res["regex"])
+        self.assertIn("f", res["regex"]["path"])
+
     def test__size_operator_not_supported(self):
         q = AtlasQ(f=3)
         t = AtlasTransform(q.query, AtlasIndex("test"))