From b70d2397dec4c2066c3de53a136f6f1db44675eb Mon Sep 17 00:00:00 2001
From: Radu Toma <radu.toma@sectorlabs.ro>
Date: Thu, 15 Oct 2020 18:05:58 +0300
Subject: [PATCH] Changed how annotations are temporarily renamed

If only a subset of annotations is renamed before calling the base class,
and then renamed back using add new key/remove old key, then the ordering of
annotations is changed (annotations are stored in a OrderedDict internally).
This will cause issues if ordering is important (e.g., in a union).
---
 psqlextra/query.py  | 24 +++++++++++-------------
 tests/test_query.py | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/psqlextra/query.py b/psqlextra/query.py
index 263ece8e..d3454499 100644
--- a/psqlextra/query.py
+++ b/psqlextra/query.py
@@ -1,3 +1,4 @@
+from collections import OrderedDict
 from itertools import chain
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
@@ -32,22 +33,19 @@ def annotate(self, **annotations):
         Normally, the annotate function doesn't allow you to use the
         name of an existing field on the model as the alias name. This
         version of the function does allow that.
-        """
-
-        fields = {field.name: field for field in self.model._meta.get_fields()}
 
-        # temporarily rename the fields that have the same
-        # name as a field name, we'll rename them back after
-        # the function in the base class ran
-        new_annotations = {}
+        This is done by temporarily renaming the fields in order to avoid the
+        check for conflicts that the base class does.
+        We rename all fields instead of the ones that already exist because
+        the annotations are stored in an OrderedDict. Renaming only the
+        conflicts will mess up the order.
+        """
+        new_annotations = OrderedDict()
         renames = {}
         for name, value in annotations.items():
-            if name in fields:
-                new_name = "%s_new" % name
-                new_annotations[new_name] = value
-                renames[new_name] = name
-            else:
-                new_annotations[name] = value
+            new_name = "%s_new" % name
+            new_annotations[new_name] = value
+            renames[new_name] = name
 
         # run the base class's annotate function
         result = super().annotate(**new_annotations)
diff --git a/tests/test_query.py b/tests/test_query.py
index b25fcf47..35d0cfe7 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -61,6 +61,20 @@ def test_query_annotate_rename_chain():
     assert obj[0]["value"] == 23
 
 
+def test_query_annotate_rename_order():
+    """Tests whether annotation order is preserved after a rename."""
+
+    model = get_fake_model(
+        {
+            "name": models.CharField(max_length=10),
+            "value": models.IntegerField(),
+        }
+    )
+
+    qs = model.objects.annotate(value=F("value"), value_2=F("value"))
+    assert list(qs.query.annotations.keys()) == ["value", "value_2"]
+
+
 def test_query_hstore_value_update_f_ref():
     """Tests whether F(..) expressions can be used in hstore values when
     performing update queries."""