From f67fd307350d8ec217deba554eea65af29410abd Mon Sep 17 00:00:00 2001 From: MohammadMahdi Khalaj Date: Mon, 6 Nov 2023 13:47:05 +0330 Subject: [PATCH 1/2] Fix `project` bug and improve `add_field` --- aggify/aggify.py | 24 +++++++++++++++++------- aggify/exceptions.py | 6 ++++++ 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/aggify/aggify.py b/aggify/aggify.py index 5422bae..d403623 100644 --- a/aggify/aggify.py +++ b/aggify/aggify.py @@ -12,6 +12,7 @@ InvalidEmbeddedField, OutStageError, InvalidArgument, + InvalidProjection, ) from aggify.types import QueryParams, CollectionType from aggify.utilty import ( @@ -84,6 +85,9 @@ def project(self, **kwargs: QueryParams) -> "Aggify": Aggify: Returns an instance of the Aggify class for potential method chaining. """ + if all([i in kwargs.values() for i in [0, 1]]): + raise InvalidProjection() + # Extract fields to keep and check if _id should be deleted to_keep_values = {"id"} delete_id = kwargs.get("id") is not None @@ -99,14 +103,16 @@ def project(self, **kwargs: QueryParams) -> "Aggify": to_keep_values.add(key) self.base_model._fields[key] = mongoengine_fields.IntField() # noqa projection[get_db_field(self.base_model, key)] = value # noqa + if value == 0: + del self.base_model._fields[key] # noqa # Remove fields from the base model, except the ones in to_keep_values and possibly _id - keys_for_deletion = self.base_model._fields.keys() - to_keep_values # noqa - if delete_id: - keys_for_deletion.add("id") - for key in keys_for_deletion: - del self.base_model._fields[key] # noqa - + if to_keep_values != {"id"}: + keys_for_deletion = self.base_model._fields.keys() - to_keep_values # noqa + if delete_id: + keys_for_deletion.add("id") + for key in keys_for_deletion: + del self.base_model._fields[key] # noqa # Append the projection stage to the pipelines self.pipelines.append({"$project": projection}) @@ -157,10 +163,14 @@ def add_fields(self, **fields) -> "Aggify": # noqa add_fields_stage["$addFields"][field] = {"$literal": expression} elif isinstance(expression, F): add_fields_stage["$addFields"][field] = expression.to_dict() - elif isinstance(expression, list): + elif isinstance(expression, (list, dict)): add_fields_stage["$addFields"][field] = expression elif isinstance(expression, Cond): add_fields_stage["$addFields"][field] = dict(expression) + elif isinstance(expression, Q): + add_fields_stage["$addFields"][field] = convert_match_query( + dict(expression) + )["$match"] else: raise AggifyValueError([str, F, list], type(expression)) # TODO: Should be checked if new field is embedded, create embedded field. diff --git a/aggify/exceptions.py b/aggify/exceptions.py index 951a123..b5ab25d 100644 --- a/aggify/exceptions.py +++ b/aggify/exceptions.py @@ -74,3 +74,9 @@ def __init__(self, expected_list: list): self.message = f"Input is not correctly passed, expected {[expected for expected in expected_list]}" self.expecteds = expected_list super().__init__(self.message) + + +class InvalidProjection(AggifyBaseException): + def __init__(self): + self.message = "You can't use inclusion and exclusion together." + super().__init__(self.message) From 4d90d31fdf434e14a22f2c9fc1d6951e21e3d4d2 Mon Sep 17 00:00:00 2001 From: MohammadMahdi Khalaj Date: Mon, 6 Nov 2023 13:47:37 +0330 Subject: [PATCH 2/2] Add tests for `project` bug and `add_field` improvement --- tests/test_aggify.py | 23 +++++++++++++++++++++-- tests/test_query.py | 4 ++-- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/tests/test_aggify.py b/tests/test_aggify.py index e49ea9f..f931deb 100644 --- a/tests/test_aggify.py +++ b/tests/test_aggify.py @@ -12,6 +12,7 @@ AlreadyExistsField, InvalidEmbeddedField, MongoIndexError, + InvalidProjection, ) @@ -53,9 +54,9 @@ def test_filtering_and_projection(self): def test_filtering_and_projection_with_deleting_id(self): aggify = Aggify(BaseModel) - aggify.filter(age__gte=30).project(name=1, age=1, id=0) + aggify.filter(age__gte=30).project(name=1, age=1, id=1) assert len(aggify.pipelines) == 2 - assert aggify.pipelines[1]["$project"] == {"_id": 0, "name": 1, "age": 1} + assert aggify.pipelines[1]["$project"] == {"_id": 1, "name": 1, "age": 1} def test_filtering_and_ordering(self): aggify = Aggify(BaseModel) @@ -617,3 +618,21 @@ def test_aggify_get_item_slice_negative_start(self): def test_group_invalid_field(self): thing = list(Aggify(BaseModel).group("invalid").annotate("name", "first", 2)) assert thing[0]["$group"] == {"_id": "invalid", "name": {"$first": 2}} + + def test_add_field_with_q(self): + thing = list(Aggify(BaseModel).add_fields(new_field=Q(name__exact="aggify"))) + assert thing[0]["$addFields"] == {"new_field": {"$eq": ["$name", "aggify"]}} + + def test_add_field_with_dict(self): + thing = list( + Aggify(BaseModel).add_fields(new_field={"$eq": ["$owner.visibility", 0]}) + ) + assert thing[0]["$addFields"] == { + "new_field": {"$eq": ["$owner.visibility", 0]} + } + + def test_project_use_inclusion_and_exclusion_together(self): + aggify = Aggify(BaseModel) + with pytest.raises(InvalidProjection): + # noinspection PyUnusedLocal + var = aggify.project(name=0, age=1) diff --git a/tests/test_query.py b/tests/test_query.py index 61ef1b2..3fe9ee7 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -82,10 +82,10 @@ class ParameterTestCase: ParameterTestCase( compiled_query=Aggify(PostDocument) .filter(caption__contains="hello") - .project(caption=1, deleted_at=0), + .project(caption=1, deleted_at=1), expected_query=[ {"$match": {"caption": {"$regex": "hello"}}}, - {"$project": {"caption": 1, "deleted_at": 0}}, + {"$project": {"caption": 1, "deleted_at": 1}}, ], ), ParameterTestCase(