From 5cde392c2ffaed1cfef755a16a8dddbc503ad953 Mon Sep 17 00:00:00 2001 From: Radu Toma Date: Thu, 28 May 2020 11:14:21 +0300 Subject: [PATCH] Added rename for annotation mask as well Django's QuerySet also stores annotation names in a 'annotation_select_mask' set. Not updating this set will cause a bug when the QuerySet gets cloned, and the renamed annotations will be lost. --- psqlextra/sql.py | 4 ++++ tests/test_query.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/psqlextra/sql.py b/psqlextra/sql.py index 216b5e82..d27094ff 100644 --- a/psqlextra/sql.py +++ b/psqlextra/sql.py @@ -55,6 +55,10 @@ def rename_annotations(self, annotations) -> None: self.annotations[new_name] = annotation del self.annotations[old_name] + if self.annotation_select_mask: + self.annotation_select_mask.discard(old_name) + self.annotation_select_mask.add(new_name) + def add_fields(self, field_names: List[str], *args, **kwargs) -> bool: """Adds the given (model) fields to the select set. diff --git a/tests/test_query.py b/tests/test_query.py index 722f3700..b25fcf47 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -43,6 +43,24 @@ def test_query_annotate_rename(): assert obj.title == "swen" +def test_query_annotate_rename_chain(): + """Tests whether annotations are behaving correctly after a QuerySet + chain.""" + + model = get_fake_model( + { + "name": models.CharField(max_length=10), + "value": models.IntegerField(), + } + ) + + model.objects.create(name="test", value=23) + + obj = model.objects.values("name").annotate(value=F("value"))[:1] + assert "value" in obj[0] + assert obj[0]["value"] == 23 + + def test_query_hstore_value_update_f_ref(): """Tests whether F(..) expressions can be used in hstore values when performing update queries."""