From 36d40f0a512d6074ecf1718b9c10dfd04d5130cc Mon Sep 17 00:00:00 2001 From: MohammadMahdi Khalaj Date: Wed, 8 Nov 2023 14:53:34 +0330 Subject: [PATCH 1/3] Add `raw_let` and `replace_root` improvements --- aggify/aggify.py | 35 ++++++++++++++++++++++++++--------- 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/aggify/aggify.py b/aggify/aggify.py index 5b16734..98a0d58 100644 --- a/aggify/aggify.py +++ b/aggify/aggify.py @@ -140,6 +140,7 @@ def order_by(self, *order_fields: Union[str, List[str]]) -> "Aggify": @last_out_stage_check def raw(self, raw_query: dict) -> "Aggify": self.pipelines.append(raw_query) + self.pipelines = self.__combine_sequential_matches() return self @last_out_stage_check @@ -558,6 +559,7 @@ def lookup( let: Union[List[str], None] = None, local_field: Union[str, None] = None, foreign_field: Union[str, None] = None, + raw_let: Union[Dict, None] = None, ) -> "Aggify": """ Generates a MongoDB lookup pipeline stage. @@ -574,6 +576,7 @@ def lookup( localField and foreignField are not used. local_field (Union[str, None], optional): The local field to join on when let not provided. foreign_field (Union[str, None], optional): The foreign field to join on when let not provided. + raw_let (Union[Dict, None]): raw let Returns: Aggify: An instance of the Aggify class representing a MongoDB lookup pipeline stage. @@ -583,11 +586,11 @@ def lookup( check_field_exists(self.base_model, as_name) # noqa from_collection_name = from_collection._meta.get("collection") # noqa - if not let and not (local_field and foreign_field): + if not (let or raw_let) and not (local_field and foreign_field): raise InvalidArgument( - expected_list=[["local_field", "foreign_field"], "let"] + expected_list=[["local_field", "foreign_field"], ["let", "raw_let"]] ) - elif not let: + elif not (let or raw_let): lookup_stage = { "$lookup": { "from": from_collection_name, @@ -602,11 +605,16 @@ def lookup( if not query: raise InvalidArgument(expected_list=["query"]) + if let is None: + let = [] + let_dict = { field: f"${get_db_field(self.base_model, self.get_field_name_recursively(field))}" # noqa for field in let } + let = list(raw_let.keys()) if let is [] else let + for q in query: # Construct the match stage for each query if isinstance(q, Q): @@ -625,6 +633,8 @@ def lookup( ) # Append the lookup stage with multiple match stages to the pipeline + if raw_let: + let_dict.update(raw_let) lookup_stage = { "$lookup": { "from": from_collection_name, @@ -675,13 +685,17 @@ def _replace_base(self, embedded_field) -> str: InvalidEmbeddedField: If the specified embedded field is not found or is not of the correct type. """ model_field = self.get_model_field(self.base_model, embedded_field) # noqa - + field_name = get_db_field(self.base_model, embedded_field) + if "__module__" in model_field.__dict__: + self.base_model._fields = ( + model_field._fields + ) # load new fields into old model + return f"${field_name}" if not hasattr(model_field, "document_type") or not issubclass( model_field.document_type, EmbeddedDocument ): raise InvalidEmbeddedField(field=embedded_field) - - return f"${model_field.db_field}" + return f"${field_name}" @last_out_stage_check def replace_root( @@ -703,10 +717,13 @@ def replace_root( """ name = self._replace_base(embedded_field) - if not merge: - new_root = {"$replaceRoot": {"$newRoot": name}} - else: + if merge: new_root = {"$replaceRoot": {"newRoot": {"$mergeObjects": [merge, name]}}} + self.base_model._fields.update( # noqa + {key: mongoengine_fields.IntField() for key, value in merge.items()} + ) + else: + new_root = {"$replaceRoot": {"$newRoot": name}} self.pipelines.append(new_root) return self From 691e037f4028e370142094f62e276ab26abde045 Mon Sep 17 00:00:00 2001 From: MohammadMahdi Khalaj Date: Wed, 8 Nov 2023 14:54:25 +0330 Subject: [PATCH 2/3] Add `raw_let` and `replace_root` tests --- tests/test_aggify.py | 18 +++++++++++++ tests/test_query.py | 62 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/tests/test_aggify.py b/tests/test_aggify.py index cc71fcf..983faf0 100644 --- a/tests/test_aggify.py +++ b/tests/test_aggify.py @@ -644,3 +644,21 @@ def test_project_add_new_field(self): thing = list(aggify.project(test="test", id=0)) assert thing[0]["$project"] == {"test": "test", "_id": 0} assert list(aggify.base_model._fields.keys()) == ["test"] + + def test_lookup_raw_let(self): + aggify = Aggify(BaseModel) + thing = list( + aggify.lookup( + BaseModel, + raw_let={"test": "$name"}, + query=[Q(name__exact="$$test")], + as_name="test_name", + ) + ) + assert thing[0]["$lookup"] == { + "from": None, + "let": {"test": "$name"}, + "pipeline": [{"$match": {"$expr": {"$eq": ["$name", "$$test"]}}}], + "as": "test_name", + } + assert "test_name" in list(aggify.base_model._fields.keys()) diff --git a/tests/test_query.py b/tests/test_query.py index 3fe9ee7..86185b1 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -466,6 +466,68 @@ class ParameterTestCase: {"$group": {"_id": "$stat.like_count", "sss": {"$first": "sss"}}} ], ), + ParameterTestCase( + compiled_query=( + Aggify(PostDocument).lookup( + AccountDocument, + let=["caption"], + raw_let={ + "latest_story_id": {"$last": {"$slice": ["$owner.story", -1]}} + }, + query=[ + Q(end__exact="caption") & Q(start__exact="$$latest_story_id._id") + ], + as_name="is_seen", + ) + ), + expected_query=[ + { + "$lookup": { + "from": "account", + "let": { + "caption": "$caption", + "latest_story_id": {"$last": {"$slice": ["$owner.story", -1]}}, + }, + "pipeline": [ + { + "$match": { + "$expr": { + "$and": [ + {"$eq": ["$end", "$$caption"]}, + {"$eq": ["$start", "$$latest_story_id._id"]}, + ] + } + } + } + ], + "as": "is_seen", + } + } + ], + ), + ParameterTestCase( + compiled_query=( + Aggify(PostDocument) + .lookup( + PostDocument, + local_field="end", + foreign_field="id", + as_name="saved_post", + ) + .replace_root(embedded_field="saved_post") + ), + expected_query=[ + { + "$lookup": { + "as": "saved_post", + "foreignField": "_id", + "from": "post_document", + "localField": "end", + } + }, + {"$replaceRoot": {"$newRoot": "$saved_post"}}, + ], + ), ] From 42e49c2841b550d7ef07ffb80db89d1901472068 Mon Sep 17 00:00:00 2001 From: MohammadMahdi Khalaj Date: Wed, 8 Nov 2023 14:56:03 +0330 Subject: [PATCH 3/3] Improve `replace_with` --- aggify/aggify.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/aggify/aggify.py b/aggify/aggify.py index 98a0d58..5712d23 100644 --- a/aggify/aggify.py +++ b/aggify/aggify.py @@ -752,6 +752,9 @@ def replace_with( new_root = {"$replaceWith": name} else: new_root = {"$replaceWith": {"$mergeObjects": [merge, name]}} + self.base_model._fields.update( # noqa + {key: mongoengine_fields.IntField() for key, value in merge.items()} + ) self.pipelines.append(new_root) return self